Jakevin commited on
Commit
7bf0e3b
·
verified ·
1 Parent(s): cef2507

Publish standalone Kev-4B ternary MLX v0.1.0 experimental

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
LICENSE ADDED
@@ -0,0 +1,202 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ Apache License
3
+ Version 2.0, January 2004
4
+ http://www.apache.org/licenses/
5
+
6
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
7
+
8
+ 1. Definitions.
9
+
10
+ "License" shall mean the terms and conditions for use, reproduction,
11
+ and distribution as defined by Sections 1 through 9 of this document.
12
+
13
+ "Licensor" shall mean the copyright owner or entity authorized by
14
+ the copyright owner that is granting the License.
15
+
16
+ "Legal Entity" shall mean the union of the acting entity and all
17
+ other entities that control, are controlled by, or are under common
18
+ control with that entity. For the purposes of this definition,
19
+ "control" means (i) the power, direct or indirect, to cause the
20
+ direction or management of such entity, whether by contract or
21
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
22
+ outstanding shares, or (iii) beneficial ownership of such entity.
23
+
24
+ "You" (or "Your") shall mean an individual or Legal Entity
25
+ exercising permissions granted by this License.
26
+
27
+ "Source" form shall mean the preferred form for making modifications,
28
+ including but not limited to software source code, documentation
29
+ source, and configuration files.
30
+
31
+ "Object" form shall mean any form resulting from mechanical
32
+ transformation or translation of a Source form, including but
33
+ not limited to compiled object code, generated documentation,
34
+ and conversions to other media types.
35
+
36
+ "Work" shall mean the work of authorship, whether in Source or
37
+ Object form, made available under the License, as indicated by a
38
+ copyright notice that is included in or attached to the work
39
+ (an example is provided in the Appendix below).
40
+
41
+ "Derivative Works" shall mean any work, whether in Source or Object
42
+ form, that is based on (or derived from) the Work and for which the
43
+ editorial revisions, annotations, elaborations, or other modifications
44
+ represent, as a whole, an original work of authorship. For the purposes
45
+ of this License, Derivative Works shall not include works that remain
46
+ separable from, or merely link (or bind by name) to the interfaces of,
47
+ the Work and Derivative Works thereof.
48
+
49
+ "Contribution" shall mean any work of authorship, including
50
+ the original version of the Work and any modifications or additions
51
+ to that Work or Derivative Works thereof, that is intentionally
52
+ submitted to Licensor for inclusion in the Work by the copyright owner
53
+ or by an individual or Legal Entity authorized to submit on behalf of
54
+ the copyright owner. For the purposes of this definition, "submitted"
55
+ means any form of electronic, verbal, or written communication sent
56
+ to the Licensor or its representatives, including but not limited to
57
+ communication on electronic mailing lists, source code control systems,
58
+ and issue tracking systems that are managed by, or on behalf of, the
59
+ Licensor for the purpose of discussing and improving the Work, but
60
+ excluding communication that is conspicuously marked or otherwise
61
+ designated in writing by the copyright owner as "Not a Contribution."
62
+
63
+ "Contributor" shall mean Licensor and any individual or Legal Entity
64
+ on behalf of whom a Contribution has been received by Licensor and
65
+ subsequently incorporated within the Work.
66
+
67
+ 2. Grant of Copyright License. Subject to the terms and conditions of
68
+ this License, each Contributor hereby grants to You a perpetual,
69
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
70
+ copyright license to reproduce, prepare Derivative Works of,
71
+ publicly display, publicly perform, sublicense, and distribute the
72
+ Work and such Derivative Works in Source or Object form.
73
+
74
+ 3. Grant of Patent License. Subject to the terms and conditions of
75
+ this License, each Contributor hereby grants to You a perpetual,
76
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
77
+ (except as stated in this section) patent license to make, have made,
78
+ use, offer to sell, sell, import, and otherwise transfer the Work,
79
+ where such license applies only to those patent claims licensable
80
+ by such Contributor that are necessarily infringed by their
81
+ Contribution(s) alone or by combination of their Contribution(s)
82
+ with the Work to which such Contribution(s) was submitted. If You
83
+ institute patent litigation against any entity (including a
84
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
85
+ or a Contribution incorporated within the Work constitutes direct
86
+ or contributory patent infringement, then any patent licenses
87
+ granted to You under this License for that Work shall terminate
88
+ as of the date such litigation is filed.
89
+
90
+ 4. Redistribution. You may reproduce and distribute copies of the
91
+ Work or Derivative Works thereof in any medium, with or without
92
+ modifications, and in Source or Object form, provided that You
93
+ meet the following conditions:
94
+
95
+ (a) You must give any other recipients of the Work or
96
+ Derivative Works a copy of this License; and
97
+
98
+ (b) You must cause any modified files to carry prominent notices
99
+ stating that You changed the files; and
100
+
101
+ (c) You must retain, in the Source form of any Derivative Works
102
+ that You distribute, all copyright, patent, trademark, and
103
+ attribution notices from the Source form of the Work,
104
+ excluding those notices that do not pertain to any part of
105
+ the Derivative Works; and
106
+
107
+ (d) If the Work includes a "NOTICE" text file as part of its
108
+ distribution, then any Derivative Works that You distribute must
109
+ include a readable copy of the attribution notices contained
110
+ within such NOTICE file, excluding those notices that do not
111
+ pertain to any part of the Derivative Works, in at least one
112
+ of the following places: within a NOTICE text file distributed
113
+ as part of the Derivative Works; within the Source form or
114
+ documentation, if provided along with the Derivative Works; or,
115
+ within a display generated by the Derivative Works, if and
116
+ wherever such third-party notices normally appear. The contents
117
+ of the NOTICE file are for informational purposes only and
118
+ do not modify the License. You may add Your own attribution
119
+ notices within Derivative Works that You distribute, alongside
120
+ or as an addendum to the NOTICE text from the Work, provided
121
+ that such additional attribution notices cannot be construed
122
+ as modifying the License.
123
+
124
+ You may add Your own copyright statement to Your modifications and
125
+ may provide additional or different license terms and conditions
126
+ for use, reproduction, or distribution of Your modifications, or
127
+ for any such Derivative Works as a whole, provided Your use,
128
+ reproduction, and distribution of the Work otherwise complies with
129
+ the conditions stated in this License.
130
+
131
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
132
+ any Contribution intentionally submitted for inclusion in the Work
133
+ by You to the Licensor shall be under the terms and conditions of
134
+ this License, without any additional terms or conditions.
135
+ Notwithstanding the above, nothing herein shall supersede or modify
136
+ the terms of any separate license agreement you may have executed
137
+ with Licensor regarding such Contributions.
138
+
139
+ 6. Trademarks. This License does not grant permission to use the trade
140
+ names, trademarks, service marks, or product names of the Licensor,
141
+ except as required for reasonable and customary use in describing the
142
+ origin of the Work and reproducing the content of the NOTICE file.
143
+
144
+ 7. Disclaimer of Warranty. Unless required by applicable law or
145
+ agreed to in writing, Licensor provides the Work (and each
146
+ Contributor provides its Contributions) on an "AS IS" BASIS,
147
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
148
+ implied, including, without limitation, any warranties or conditions
149
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
150
+ PARTICULAR PURPOSE. You are solely responsible for determining the
151
+ appropriateness of using or redistributing the Work and assume any
152
+ risks associated with Your exercise of permissions under this License.
153
+
154
+ 8. Limitation of Liability. In no event and under no legal theory,
155
+ whether in tort (including negligence), contract, or otherwise,
156
+ unless required by applicable law (such as deliberate and grossly
157
+ negligent acts) or agreed to in writing, shall any Contributor be
158
+ liable to You for damages, including any direct, indirect, special,
159
+ incidental, or consequential damages of any character arising as a
160
+ result of this License or out of the use or inability to use the
161
+ Work (including but not limited to damages for loss of goodwill,
162
+ work stoppage, computer failure or malfunction, or any and all
163
+ other commercial damages or losses), even if such Contributor
164
+ has been advised of the possibility of such damages.
165
+
166
+ 9. Accepting Warranty or Additional Liability. While redistributing
167
+ the Work or Derivative Works thereof, You may choose to offer,
168
+ and charge a fee for, acceptance of support, warranty, indemnity,
169
+ or other liability obligations and/or rights consistent with this
170
+ License. However, in accepting such obligations, You may act only
171
+ on Your own behalf and on Your sole responsibility, not on behalf
172
+ of any other Contributor, and only if You agree to indemnify,
173
+ defend, and hold each Contributor harmless for any liability
174
+ incurred by, or claims asserted against, such Contributor by reason
175
+ of your accepting any such warranty or additional liability.
176
+
177
+ END OF TERMS AND CONDITIONS
178
+
179
+ APPENDIX: How to apply the Apache License to your work.
180
+
181
+ To apply the Apache License to your work, attach the following
182
+ boilerplate notice, with the fields enclosed by brackets "[]"
183
+ replaced with your own identifying information. (Don't include
184
+ the brackets!) The text should be enclosed in the appropriate
185
+ comment syntax for the file format. We also recommend that a
186
+ file or class name and description of purpose be included on the
187
+ same "printed page" as the copyright notice for easier
188
+ identification within third-party archives.
189
+
190
+ Copyright 2026 Jared Palmer
191
+
192
+ Licensed under the Apache License, Version 2.0 (the "License");
193
+ you may not use this file except in compliance with the License.
194
+ You may obtain a copy of the License at
195
+
196
+ http://www.apache.org/licenses/LICENSE-2.0
197
+
198
+ Unless required by applicable law or agreed to in writing, software
199
+ distributed under the License is distributed on an "AS IS" BASIS,
200
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
201
+ See the License for the specific language governing permissions and
202
+ limitations under the License.
NOTICE ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ Kev source and original decision model: Copyright 2026 Jared Palmer. Apache-2.0.
2
+ Base backbone: Qwen/Qwen3.5-4B-Base, Apache-2.0.
3
+ Standalone mixed ternary/4-bit compression export and packaging: Jakevin, 2026-10-06.
4
+ This distribution modifies kev/checkpoint.py and kev/mlx_model.py and adds kev/packed_model.py.
5
+ Upstream source: https://github.com/jaredpalmer/kev
6
+ Upstream Kev checkpoint: jaredpalmer/kev-4b@6cfce5c2fa4b4bd64026336ab649c5ca78857d52
7
+ Base checkpoint: Qwen/Qwen3.5-4B-Base@1001bb4d826a52d1f399e183466143f4da7b741b
README.md ADDED
@@ -0,0 +1,176 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ base_model: jaredpalmer/kev-4b
4
+ base_model_relation: quantized
5
+ library_name: mlx
6
+ tags:
7
+ - kev
8
+ - decision-model
9
+ - ternary
10
+ - mixed-precision
11
+ - post-training-quantization
12
+ - mlx
13
+ - experimental
14
+ ---
15
+
16
+ # Kev-4B Ternary + CLoQ16 for MLX
17
+
18
+ An experimental, standalone compression of **Jared Palmer's Kev-4B decision
19
+ model**, published by Jakevin. The entire distribution is approximately
20
+ **1.71 GB**, including the compressed backbone, compensation adapter, pointer
21
+ head, configuration, tokenizer, and the required Kev runtime source.
22
+
23
+ On **decision-v7 development**, 1,264 scored clean questions, native packed
24
+ inference reaches **84.41% accuracy**, retaining **96.82%** of the archived
25
+ bf16-MLX accuracy. These are development diagnostics for a suite related to
26
+ Kev's training recipe, not a fresh blind test or a general retention guarantee.
27
+
28
+ This is a **decision scorer**, not a text-generating chat model. It implements
29
+ Kev's existing `state` / `questions` -> `answers` / `usage` System One API.
30
+ **Apple Silicon and the included custom MLX runtime are required.** Stock
31
+ Transformers, stock mlx-lm text generation, and llama.cpp are not supported
32
+ loading paths for this artifact.
33
+
34
+ ## Download and Run
35
+
36
+ Download this repository, including `runtime/`. It contains the packed-aware
37
+ Kev loader and serving source; the upstream Kev release alone does not include
38
+ this loader. Python 3.12 or 3.13 is required.
39
+
40
+ ```sh
41
+ python3 -m pip install huggingface_hub
42
+ python3 -c "from huggingface_hub import snapshot_download; snapshot_download('Jakevin/kev-4b-ternary-mlx', revision='v0.1.0-experimental', local_dir='kev-4b-ternary-mlx')"
43
+ cd kev-4b-ternary-mlx
44
+ python3.13 -m venv .venv
45
+ .venv/bin/python -m pip install './runtime[serve]'
46
+
47
+ HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-hf-cache KEV_BACKEND=mlx \
48
+ KEV_PREFIX_CACHE=1 KEV_PREFIX_MAX_TOKENS=4096 \
49
+ .venv/bin/python -m kev.serve --run . --host 127.0.0.1 --port 8010
50
+ ```
51
+
52
+ ```sh
53
+ curl http://127.0.0.1:8010/v1/systemone \
54
+ -H 'Content-Type: application/json' \
55
+ -d '{"state":"I was charged twice for order 8812.","questions":{"team":{"type":"choice","instructions":"Which team should handle this ticket?","criteria":{"billing":"Charges and payments","shipping":"Delivery and delays"}}}}'
56
+ ```
57
+
58
+ The library interface is unchanged:
59
+
60
+ ```python
61
+ from kev.checkpoint import Checkpoint
62
+
63
+ tokenizer, model = Checkpoint("path/to/downloaded-repository").load("mps")
64
+ ```
65
+
66
+ After dependencies are installed, loading and inference work offline with an
67
+ empty Hugging Face cache. They do not load the original backbone or original
68
+ PEFT adapter, reconstruct a dense bf16 backbone, or construct the dense
69
+ compensation matrix. Architecture initialization is lazy and its random dense
70
+ parameters are replaced before evaluation.
71
+
72
+ ## Compression Recipe
73
+
74
+ Source checkpoint: [`jaredpalmer/kev-4b`](https://huggingface.co/jaredpalmer/kev-4b),
75
+ revision `6cfce5c2fa4b4bd64026336ab649c5ca78857d52`, based on
76
+ [`Qwen/Qwen3.5-4B-Base`](https://huggingface.co/Qwen/Qwen3.5-4B-Base), revision
77
+ `1001bb4d826a52d1f399e183466143f4da7b741b`.
78
+
79
+ - The source Kev LoRA was merged before the original compression experiment.
80
+ - GPTQ-style error compensation uses 128 records from the decision-v7
81
+ **calibration** partition, RNG 1234, damping 0.01 and sequential 2-layer chunks.
82
+ - Most backbone linear weights are ternary, with an fp16 scale per 64 input columns.
83
+ - Gated DeltaNet `out_proj` and full-attention `q/k/v/o_proj` use 4-bit affine
84
+ weights with group size 64. The embedding is also 4-bit.
85
+ - Small `in_proj_a/b`, norms, conv kernels and DeltaNet parameters retain their
86
+ saved source precision. The pointer head and source temperature are unchanged.
87
+ - CLoQ-style rank-16 compensation uses a closed-form calibration-based solution;
88
+ no gradient retraining or new temperature fitting was performed.
89
+ - No online rotation is required by this recipe.
90
+
91
+ At inference, ternary codes use native **2-bit** MLX affine storage. Native
92
+ storage is larger than the original five-trits-per-byte archival format. The
93
+ CLoQ factors remain bf16 and execute as two separate small matrix products
94
+ with fp32 accumulation. Activations are bf16. Only requested embedding rows
95
+ are dequantized. This is a weight compression method; activations and the
96
+ attention/prefix cache are not claimed to be 2-bit or 4-bit.
97
+
98
+ The method combines existing quantization/compensation ideas in a layer-specific
99
+ recipe; it does not claim a new GPTQ or CLoQ algorithm. Relevant work:
100
+ [GPTQ](https://arxiv.org/abs/2210.17323),
101
+ [CLoQ](https://arxiv.org/abs/2501.18475).
102
+
103
+ ## Development Evaluation
104
+
105
+ All 1,204 decision-v7 development records were evaluated: 1,468 total rows,
106
+ 1,264 scored clean questions, zero rejected records and zero truncations.
107
+ The source checkpoint's temperature was retained.
108
+
109
+ | Path | Backbone + compensation | Accuracy | Relative accuracy retention | Brier |
110
+ |---|---:|---:|---:|---:|
111
+ | Archived bf16-MLX baseline | 8.412 GB | 0.87184 | 100% | 0.18248 |
112
+ | Archived five-trits-per-byte artifact, decoded/merged to bf16 for scoring | 1.434 GB | 0.84335 | 96.73% | 0.23260 |
113
+ | This native MLX packed inference path | 1.679 GB | **0.84415** | **96.82%** | **0.23251** |
114
+
115
+ The complete published folder is approximately 1.71 GB; sizes above exclude
116
+ tokenizer and pointer head. GB means decimal GB.
117
+
118
+ Native quantized matmul plus separate compensation has different arithmetic
119
+ from the archived merged-bf16 simulation. Across 1,468 rows, 1 argmax changed;
120
+ mean per-row maximum probability difference was 0.00268, maximum 0.06209.
121
+ This is not bit-identical serving. The original bf16 model was not loaded again
122
+ for this comparison; its metrics came from the frozen local experiment.
123
+
124
+ Aggregate evaluation and memory records are in `evaluation/`. No evaluation
125
+ sentences, labels or per-example prediction dumps are uploaded. No NAER /
126
+ Chinese Linguipedia data, derivatives or predictions were used in this export
127
+ or included in this repository.
128
+
129
+ ## Memory Measurements
130
+
131
+ Fresh-process probes on an **Apple M4 Pro, 64 GiB RAM**, macOS 26.6.2:
132
+ MLX 0.32.2, mlx-lm 0.31.3, torch 2.8.0, transformers 5.17.0.
133
+ Original loaders and downloads were explicitly blocked during the probes.
134
+
135
+ | Measurement | Short request | Longer request |
136
+ |---|---:|---:|
137
+ | State / total input tokens | 19 / 42 | 3,519 / 3,542 |
138
+ | Load time | 0.57 s | 0.57 s |
139
+ | MLX active weights after load | 1.679 GB | 1.679 GB |
140
+ | Process physical-footprint peak while loading | 2.396 GB | 2.398 GB |
141
+ | MLX inference peak | 2.013 GB | 3.066 GB |
142
+ | Process physical-footprint inference peak | **2.846 GB** | **3.862 GB** |
143
+ | First inference time | 0.155 s | 6.599 s |
144
+
145
+ Physical footprint includes Metal buffers; RSS alone undercounts them. These
146
+ are single-request probes, not minimum-RAM device certification. **No 8-GB
147
+ hardware was tested.** More state tokens, questions, or retained prefixes can
148
+ increase memory. The synthetic longer probe measures memory/latency, not
149
+ long-context accuracy. `KEV_PREFIX_MAX_TOKENS` limits retained prefix cache,
150
+ not input admission. `mx.set_memory_limit(3 GB)` is a guideline rather than a
151
+ hard RAM cap; the longer probe exceeded it.
152
+
153
+ ## Reproducibility and Limitations
154
+
155
+ - The runtime and the native exporter are in `runtime/`. Source and model
156
+ hashes are recorded in `packed.json`, `head.pt["compression"]` and
157
+ `release-manifest.json`.
158
+ - The test source is included. The local release check passed 7 relevant
159
+ tests, including native code/adapter arithmetic, loading with base loaders
160
+ blocked, missing-weight rejection, and the System One API.
161
+ - Decision-v7 development was used to select compression recipes. No new
162
+ blind benchmark, transfer-v9, hard-v1 or breadth-v1 evaluation is claimed.
163
+ - Full 3-bit and 4-bit same-condition baselines were not evaluated here.
164
+ - The source temperature is retained; Brier/ECE worsened relative to the
165
+ bf16 baseline. Confidence-sensitive users should evaluate their own task.
166
+ - Native inference has not been validated on non-Apple devices or llama.cpp.
167
+ - This is an experimental derivative of Kev-4B, not a new independently
168
+ pretrained foundation model and not the Taiwan/mainland wording model.
169
+
170
+ ## Attribution and License
171
+
172
+ Apache-2.0. The original Kev decision model and source are by **Jared Palmer**;
173
+ the base backbone is from **Qwen**. This compression experiment, standalone
174
+ export and publication are by **Jakevin**. The runtime distribution modifies
175
+ Kev's checkpoint and MLX loaders and adds the packed model implementation.
176
+ See `LICENSE` and `NOTICE` for attribution.
added_tokens.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "</think>": 151668,
3
+ "</tool_call>": 151658,
4
+ "</tool_response>": 151666,
5
+ "<think>": 151667,
6
+ "<tool_call>": 151657,
7
+ "<tool_response>": 151665,
8
+ "<|box_end|>": 151649,
9
+ "<|box_start|>": 151648,
10
+ "<|endoftext|>": 151643,
11
+ "<|file_sep|>": 151664,
12
+ "<|fim_middle|>": 151660,
13
+ "<|fim_pad|>": 151662,
14
+ "<|fim_prefix|>": 151659,
15
+ "<|fim_suffix|>": 151661,
16
+ "<|im_end|>": 151645,
17
+ "<|im_start|>": 151644,
18
+ "<|image_pad|>": 151655,
19
+ "<|object_ref_end|>": 151647,
20
+ "<|object_ref_start|>": 151646,
21
+ "<|quad_end|>": 151651,
22
+ "<|quad_start|>": 151650,
23
+ "<|repo_name|>": 151663,
24
+ "<|video_pad|>": 151656,
25
+ "<|vision_end|>": 151653,
26
+ "<|vision_pad|>": 151654,
27
+ "<|vision_start|>": 151652
28
+ }
config.json ADDED
@@ -0,0 +1,104 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Qwen3_5ForConditionalGeneration"
4
+ ],
5
+ "image_token_id": 248056,
6
+ "model_type": "qwen3_5",
7
+ "text_config": {
8
+ "attention_bias": false,
9
+ "attention_dropout": 0.0,
10
+ "attn_output_gate": true,
11
+ "dtype": "bfloat16",
12
+ "eos_token_id": 248044,
13
+ "full_attention_interval": 4,
14
+ "head_dim": 256,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2560,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 9216,
19
+ "layer_types": [
20
+ "linear_attention",
21
+ "linear_attention",
22
+ "linear_attention",
23
+ "full_attention",
24
+ "linear_attention",
25
+ "linear_attention",
26
+ "linear_attention",
27
+ "full_attention",
28
+ "linear_attention",
29
+ "linear_attention",
30
+ "linear_attention",
31
+ "full_attention",
32
+ "linear_attention",
33
+ "linear_attention",
34
+ "linear_attention",
35
+ "full_attention",
36
+ "linear_attention",
37
+ "linear_attention",
38
+ "linear_attention",
39
+ "full_attention",
40
+ "linear_attention",
41
+ "linear_attention",
42
+ "linear_attention",
43
+ "full_attention",
44
+ "linear_attention",
45
+ "linear_attention",
46
+ "linear_attention",
47
+ "full_attention",
48
+ "linear_attention",
49
+ "linear_attention",
50
+ "linear_attention",
51
+ "full_attention"
52
+ ],
53
+ "linear_conv_kernel_dim": 4,
54
+ "linear_key_head_dim": 128,
55
+ "linear_num_key_heads": 16,
56
+ "linear_num_value_heads": 32,
57
+ "linear_value_head_dim": 128,
58
+ "max_position_embeddings": 262144,
59
+ "mlp_only_layers": [],
60
+ "model_type": "qwen3_5_text",
61
+ "mtp_num_hidden_layers": 1,
62
+ "mtp_use_dedicated_embeddings": false,
63
+ "num_attention_heads": 16,
64
+ "num_hidden_layers": 32,
65
+ "num_key_value_heads": 4,
66
+ "rms_norm_eps": 1e-06,
67
+ "tie_word_embeddings": true,
68
+ "use_cache": true,
69
+ "vocab_size": 248320,
70
+ "mamba_ssm_dtype": "float32",
71
+ "rope_parameters": {
72
+ "mrope_interleaved": true,
73
+ "mrope_section": [
74
+ 11,
75
+ 11,
76
+ 10
77
+ ],
78
+ "rope_type": "default",
79
+ "rope_theta": 10000000,
80
+ "partial_rotary_factor": 0.25
81
+ }
82
+ },
83
+ "tie_word_embeddings": true,
84
+ "transformers_version": "4.57.0.dev0",
85
+ "video_token_id": 248057,
86
+ "vision_config": {
87
+ "deepstack_visual_indexes": [],
88
+ "depth": 24,
89
+ "hidden_act": "gelu_pytorch_tanh",
90
+ "hidden_size": 1024,
91
+ "in_channels": 3,
92
+ "initializer_range": 0.02,
93
+ "intermediate_size": 4096,
94
+ "model_type": "qwen3_5",
95
+ "num_heads": 16,
96
+ "num_position_embeddings": 2304,
97
+ "out_hidden_size": 2560,
98
+ "patch_size": 16,
99
+ "spatial_merge_size": 2,
100
+ "temporal_patch_size": 2
101
+ },
102
+ "vision_end_token_id": 248054,
103
+ "vision_start_token_id": 248053
104
+ }
evaluation/development-summary.json ADDED
@@ -0,0 +1,91 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "clean": {
3
+ "n": 1264,
4
+ "nll": 0.43551092937105246,
5
+ "acc": 0.8441455696202531,
6
+ "ece": 0.043164737775791336,
7
+ "brier": 0.23251480498478275,
8
+ "mean_conf": 0.8088378144123675,
9
+ "confident_error_rate": 0.010284810126582278,
10
+ "coverage_at_0_9": 0.4849683544303797,
11
+ "accuracy_at_0_9": 0.9787928221859706,
12
+ "coverage_at_5pct_error": 0.6724683544303798,
13
+ "coverage_at_1pct_error": 0.16772151898734178,
14
+ "aurc": 0.046630371957041034,
15
+ "error_rate_at_0_9": 0.021207177814029365,
16
+ "confidence_bias": -0.03530775520788565,
17
+ "top_bins": {
18
+ "0.9": {
19
+ "n": 613,
20
+ "errors": 13,
21
+ "error_rate": 0.021207177814029365
22
+ },
23
+ "0.95": {
24
+ "n": 443,
25
+ "errors": 6,
26
+ "error_rate": 0.013544018058690745
27
+ },
28
+ "0.99": {
29
+ "n": 152,
30
+ "errors": 2,
31
+ "error_rate": 0.013157894736842105
32
+ }
33
+ },
34
+ "selective": {
35
+ "0.5": {
36
+ "coverage": 0.5,
37
+ "accuracy": 0.9778481012658228,
38
+ "confidence_cutoff": 0.8935324549674988
39
+ },
40
+ "0.8": {
41
+ "coverage": 0.8006329113924051,
42
+ "accuracy": 0.9061264822134387,
43
+ "confidence_cutoff": 0.5975180090213416
44
+ }
45
+ },
46
+ "score_mae": 0.5421033783798573,
47
+ "ranked_probability_score": 0.0842686503245251
48
+ },
49
+ "coverage": {
50
+ "requested_records": 1204,
51
+ "requested_questions": 1468,
52
+ "evaluated_records": 1204,
53
+ "evaluated_questions": 1468,
54
+ "rejected_records": 0,
55
+ "truncated_records": 0
56
+ },
57
+ "latency_ms": {
58
+ "median": 304.7052710317075,
59
+ "p95": 1490.851858188398
60
+ },
61
+ "calibration": {
62
+ "inference_temperature": 2.406050072164233,
63
+ "additional_temperature": 1.0,
64
+ "logits_recorded": true
65
+ },
66
+ "environment": {
67
+ "packages": {
68
+ "torch": "2.8.0",
69
+ "transformers": "5.17.0",
70
+ "peft": "0.21.0",
71
+ "flash-linear-attention": null,
72
+ "triton": null,
73
+ "causal-conv1d": null,
74
+ "mlx": "0.32.2",
75
+ "mlx-lm": "0.31.3"
76
+ },
77
+ "device": "mps",
78
+ "gpu": null,
79
+ "backend": "mlx",
80
+ "dtype": "bfloat16",
81
+ "triton_f32_default": null,
82
+ "attention": null,
83
+ "deltanet": null
84
+ },
85
+ "suite": "decision-v7",
86
+ "partition": "development",
87
+ "role": "in-distribution development diagnostic, not a blind test",
88
+ "source_baseline": "archived bf16-MLX metrics; original weights not loaded for this comparison",
89
+ "baseline_accuracy": 0.8718354430379747,
90
+ "accuracy_retention": 0.9682395644283122
91
+ }
evaluation/memory-long.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "mlx",
3
+ "dtype": "bfloat16",
4
+ "platform": "macOS-26.6.2-arm64-arm-64bit-Mach-O",
5
+ "original_loaders_blocked": true,
6
+ "mlx_limit_bytes": 3000000000,
7
+ "load_seconds": 0.5740630407817662,
8
+ "mlx_load_peak_bytes": 1678767640,
9
+ "mlx_after_load_bytes": 1678767624,
10
+ "load_footprint_peak_bytes": 2397899416,
11
+ "load_rss_peak_bytes": 2381381632,
12
+ "state_tokens": 3519,
13
+ "input_tokens": 3542,
14
+ "inference_seconds": 6.599218292161822,
15
+ "probabilities": {
16
+ "billing": 0.941597580909729,
17
+ "shipping": 0.05840245634317398
18
+ },
19
+ "mlx_inference_peak_bytes": 3065702282,
20
+ "inference_footprint_peak_bytes": 3862170744,
21
+ "inference_rss_peak_bytes": 2520678400,
22
+ "mlx_limit_is_hard_cap": false,
23
+ "machine": "Apple M4 Pro",
24
+ "machine_ram_bytes": 68719476736,
25
+ "versions": {
26
+ "mlx": "0.32.2",
27
+ "mlx-lm": "0.31.3",
28
+ "torch": "2.8.0",
29
+ "transformers": "5.17.0"
30
+ }
31
+ }
evaluation/memory-short.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "mlx",
3
+ "dtype": "bfloat16",
4
+ "platform": "macOS-26.6.2-arm64-arm-64bit-Mach-O",
5
+ "original_loaders_blocked": true,
6
+ "mlx_limit_bytes": 3000000000,
7
+ "load_seconds": 0.5707792090252042,
8
+ "mlx_load_peak_bytes": 1678767640,
9
+ "mlx_after_load_bytes": 1678767624,
10
+ "load_footprint_peak_bytes": 2396031616,
11
+ "load_rss_peak_bytes": 2307145728,
12
+ "state_tokens": 19,
13
+ "input_tokens": 42,
14
+ "inference_seconds": 0.15510379197075963,
15
+ "probabilities": {
16
+ "billing": 0.9229481220245361,
17
+ "shipping": 0.07705190777778625
18
+ },
19
+ "mlx_inference_peak_bytes": 2012569374,
20
+ "inference_footprint_peak_bytes": 2846362528,
21
+ "inference_rss_peak_bytes": 2516697088,
22
+ "mlx_limit_is_hard_cap": false,
23
+ "machine": "Apple M4 Pro",
24
+ "machine_ram_bytes": 68719476736,
25
+ "versions": {
26
+ "mlx": "0.32.2",
27
+ "mlx-lm": "0.31.3",
28
+ "torch": "2.8.0",
29
+ "transformers": "5.17.0"
30
+ }
31
+ }
evaluation/native-comparison.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "rows": 1468,
3
+ "max_abs_dp": 0.06208749928657609,
4
+ "mean_max_abs_dp": 0.0026799194979893045,
5
+ "argmax_flips": 1,
6
+ "accuracy": 0.8441455696202531,
7
+ "retention": 0.9682395644283122,
8
+ "reference": "archival packed merged-bf16 evaluation; no original weights loaded"
9
+ }
head.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6f6740d94a49187770b68c0f058d85bedc0cf7efc20dff7ca9a418c8b461dc0e
3
+ size 5247807
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c4899eaa9fb206e77260f832d00188c9b90750724422ea0eea43ed69faeff43d
3
+ size 1678909327
packed.json ADDED
@@ -0,0 +1,1016 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "format": "kev-packed-mlx-v1",
3
+ "linears": {
4
+ "language_model.model.layers.0.mlp.up_proj": {
5
+ "bits": 2,
6
+ "group_size": 64,
7
+ "rank": 16
8
+ },
9
+ "language_model.model.layers.0.mlp.down_proj": {
10
+ "bits": 2,
11
+ "group_size": 64,
12
+ "rank": 16
13
+ },
14
+ "language_model.model.layers.0.mlp.gate_proj": {
15
+ "bits": 2,
16
+ "group_size": 64,
17
+ "rank": 16
18
+ },
19
+ "language_model.model.layers.0.linear_attn.out_proj": {
20
+ "bits": 4,
21
+ "group_size": 64,
22
+ "rank": 16
23
+ },
24
+ "language_model.model.layers.0.linear_attn.in_proj_z": {
25
+ "bits": 2,
26
+ "group_size": 64,
27
+ "rank": 16
28
+ },
29
+ "language_model.model.layers.0.linear_attn.in_proj_qkv": {
30
+ "bits": 2,
31
+ "group_size": 64,
32
+ "rank": 16
33
+ },
34
+ "language_model.model.layers.1.mlp.up_proj": {
35
+ "bits": 2,
36
+ "group_size": 64,
37
+ "rank": 16
38
+ },
39
+ "language_model.model.layers.1.mlp.down_proj": {
40
+ "bits": 2,
41
+ "group_size": 64,
42
+ "rank": 16
43
+ },
44
+ "language_model.model.layers.1.mlp.gate_proj": {
45
+ "bits": 2,
46
+ "group_size": 64,
47
+ "rank": 16
48
+ },
49
+ "language_model.model.layers.1.linear_attn.out_proj": {
50
+ "bits": 4,
51
+ "group_size": 64,
52
+ "rank": 16
53
+ },
54
+ "language_model.model.layers.1.linear_attn.in_proj_z": {
55
+ "bits": 2,
56
+ "group_size": 64,
57
+ "rank": 16
58
+ },
59
+ "language_model.model.layers.1.linear_attn.in_proj_qkv": {
60
+ "bits": 2,
61
+ "group_size": 64,
62
+ "rank": 16
63
+ },
64
+ "language_model.model.layers.2.mlp.up_proj": {
65
+ "bits": 2,
66
+ "group_size": 64,
67
+ "rank": 16
68
+ },
69
+ "language_model.model.layers.2.mlp.down_proj": {
70
+ "bits": 2,
71
+ "group_size": 64,
72
+ "rank": 16
73
+ },
74
+ "language_model.model.layers.2.mlp.gate_proj": {
75
+ "bits": 2,
76
+ "group_size": 64,
77
+ "rank": 16
78
+ },
79
+ "language_model.model.layers.2.linear_attn.out_proj": {
80
+ "bits": 4,
81
+ "group_size": 64,
82
+ "rank": 16
83
+ },
84
+ "language_model.model.layers.2.linear_attn.in_proj_z": {
85
+ "bits": 2,
86
+ "group_size": 64,
87
+ "rank": 16
88
+ },
89
+ "language_model.model.layers.2.linear_attn.in_proj_qkv": {
90
+ "bits": 2,
91
+ "group_size": 64,
92
+ "rank": 16
93
+ },
94
+ "language_model.model.layers.3.mlp.up_proj": {
95
+ "bits": 2,
96
+ "group_size": 64,
97
+ "rank": 16
98
+ },
99
+ "language_model.model.layers.3.mlp.down_proj": {
100
+ "bits": 2,
101
+ "group_size": 64,
102
+ "rank": 16
103
+ },
104
+ "language_model.model.layers.3.mlp.gate_proj": {
105
+ "bits": 2,
106
+ "group_size": 64,
107
+ "rank": 16
108
+ },
109
+ "language_model.model.layers.3.self_attn.o_proj": {
110
+ "bits": 4,
111
+ "group_size": 64,
112
+ "rank": 16
113
+ },
114
+ "language_model.model.layers.3.self_attn.v_proj": {
115
+ "bits": 4,
116
+ "group_size": 64,
117
+ "rank": 16
118
+ },
119
+ "language_model.model.layers.3.self_attn.k_proj": {
120
+ "bits": 4,
121
+ "group_size": 64,
122
+ "rank": 16
123
+ },
124
+ "language_model.model.layers.3.self_attn.q_proj": {
125
+ "bits": 4,
126
+ "group_size": 64,
127
+ "rank": 16
128
+ },
129
+ "language_model.model.layers.4.mlp.up_proj": {
130
+ "bits": 2,
131
+ "group_size": 64,
132
+ "rank": 16
133
+ },
134
+ "language_model.model.layers.4.mlp.down_proj": {
135
+ "bits": 2,
136
+ "group_size": 64,
137
+ "rank": 16
138
+ },
139
+ "language_model.model.layers.4.mlp.gate_proj": {
140
+ "bits": 2,
141
+ "group_size": 64,
142
+ "rank": 16
143
+ },
144
+ "language_model.model.layers.4.linear_attn.out_proj": {
145
+ "bits": 4,
146
+ "group_size": 64,
147
+ "rank": 16
148
+ },
149
+ "language_model.model.layers.4.linear_attn.in_proj_z": {
150
+ "bits": 2,
151
+ "group_size": 64,
152
+ "rank": 16
153
+ },
154
+ "language_model.model.layers.4.linear_attn.in_proj_qkv": {
155
+ "bits": 2,
156
+ "group_size": 64,
157
+ "rank": 16
158
+ },
159
+ "language_model.model.layers.5.mlp.up_proj": {
160
+ "bits": 2,
161
+ "group_size": 64,
162
+ "rank": 16
163
+ },
164
+ "language_model.model.layers.5.mlp.down_proj": {
165
+ "bits": 2,
166
+ "group_size": 64,
167
+ "rank": 16
168
+ },
169
+ "language_model.model.layers.5.mlp.gate_proj": {
170
+ "bits": 2,
171
+ "group_size": 64,
172
+ "rank": 16
173
+ },
174
+ "language_model.model.layers.5.linear_attn.out_proj": {
175
+ "bits": 4,
176
+ "group_size": 64,
177
+ "rank": 16
178
+ },
179
+ "language_model.model.layers.5.linear_attn.in_proj_z": {
180
+ "bits": 2,
181
+ "group_size": 64,
182
+ "rank": 16
183
+ },
184
+ "language_model.model.layers.5.linear_attn.in_proj_qkv": {
185
+ "bits": 2,
186
+ "group_size": 64,
187
+ "rank": 16
188
+ },
189
+ "language_model.model.layers.6.mlp.up_proj": {
190
+ "bits": 2,
191
+ "group_size": 64,
192
+ "rank": 16
193
+ },
194
+ "language_model.model.layers.6.mlp.down_proj": {
195
+ "bits": 2,
196
+ "group_size": 64,
197
+ "rank": 16
198
+ },
199
+ "language_model.model.layers.6.mlp.gate_proj": {
200
+ "bits": 2,
201
+ "group_size": 64,
202
+ "rank": 16
203
+ },
204
+ "language_model.model.layers.6.linear_attn.out_proj": {
205
+ "bits": 4,
206
+ "group_size": 64,
207
+ "rank": 16
208
+ },
209
+ "language_model.model.layers.6.linear_attn.in_proj_z": {
210
+ "bits": 2,
211
+ "group_size": 64,
212
+ "rank": 16
213
+ },
214
+ "language_model.model.layers.6.linear_attn.in_proj_qkv": {
215
+ "bits": 2,
216
+ "group_size": 64,
217
+ "rank": 16
218
+ },
219
+ "language_model.model.layers.7.mlp.up_proj": {
220
+ "bits": 2,
221
+ "group_size": 64,
222
+ "rank": 16
223
+ },
224
+ "language_model.model.layers.7.mlp.down_proj": {
225
+ "bits": 2,
226
+ "group_size": 64,
227
+ "rank": 16
228
+ },
229
+ "language_model.model.layers.7.mlp.gate_proj": {
230
+ "bits": 2,
231
+ "group_size": 64,
232
+ "rank": 16
233
+ },
234
+ "language_model.model.layers.7.self_attn.o_proj": {
235
+ "bits": 4,
236
+ "group_size": 64,
237
+ "rank": 16
238
+ },
239
+ "language_model.model.layers.7.self_attn.v_proj": {
240
+ "bits": 4,
241
+ "group_size": 64,
242
+ "rank": 16
243
+ },
244
+ "language_model.model.layers.7.self_attn.k_proj": {
245
+ "bits": 4,
246
+ "group_size": 64,
247
+ "rank": 16
248
+ },
249
+ "language_model.model.layers.7.self_attn.q_proj": {
250
+ "bits": 4,
251
+ "group_size": 64,
252
+ "rank": 16
253
+ },
254
+ "language_model.model.layers.8.mlp.up_proj": {
255
+ "bits": 2,
256
+ "group_size": 64,
257
+ "rank": 16
258
+ },
259
+ "language_model.model.layers.8.mlp.down_proj": {
260
+ "bits": 2,
261
+ "group_size": 64,
262
+ "rank": 16
263
+ },
264
+ "language_model.model.layers.8.mlp.gate_proj": {
265
+ "bits": 2,
266
+ "group_size": 64,
267
+ "rank": 16
268
+ },
269
+ "language_model.model.layers.8.linear_attn.out_proj": {
270
+ "bits": 4,
271
+ "group_size": 64,
272
+ "rank": 16
273
+ },
274
+ "language_model.model.layers.8.linear_attn.in_proj_z": {
275
+ "bits": 2,
276
+ "group_size": 64,
277
+ "rank": 16
278
+ },
279
+ "language_model.model.layers.8.linear_attn.in_proj_qkv": {
280
+ "bits": 2,
281
+ "group_size": 64,
282
+ "rank": 16
283
+ },
284
+ "language_model.model.layers.9.mlp.up_proj": {
285
+ "bits": 2,
286
+ "group_size": 64,
287
+ "rank": 16
288
+ },
289
+ "language_model.model.layers.9.mlp.down_proj": {
290
+ "bits": 2,
291
+ "group_size": 64,
292
+ "rank": 16
293
+ },
294
+ "language_model.model.layers.9.mlp.gate_proj": {
295
+ "bits": 2,
296
+ "group_size": 64,
297
+ "rank": 16
298
+ },
299
+ "language_model.model.layers.9.linear_attn.out_proj": {
300
+ "bits": 4,
301
+ "group_size": 64,
302
+ "rank": 16
303
+ },
304
+ "language_model.model.layers.9.linear_attn.in_proj_z": {
305
+ "bits": 2,
306
+ "group_size": 64,
307
+ "rank": 16
308
+ },
309
+ "language_model.model.layers.9.linear_attn.in_proj_qkv": {
310
+ "bits": 2,
311
+ "group_size": 64,
312
+ "rank": 16
313
+ },
314
+ "language_model.model.layers.10.mlp.up_proj": {
315
+ "bits": 2,
316
+ "group_size": 64,
317
+ "rank": 16
318
+ },
319
+ "language_model.model.layers.10.mlp.down_proj": {
320
+ "bits": 2,
321
+ "group_size": 64,
322
+ "rank": 16
323
+ },
324
+ "language_model.model.layers.10.mlp.gate_proj": {
325
+ "bits": 2,
326
+ "group_size": 64,
327
+ "rank": 16
328
+ },
329
+ "language_model.model.layers.10.linear_attn.out_proj": {
330
+ "bits": 4,
331
+ "group_size": 64,
332
+ "rank": 16
333
+ },
334
+ "language_model.model.layers.10.linear_attn.in_proj_z": {
335
+ "bits": 2,
336
+ "group_size": 64,
337
+ "rank": 16
338
+ },
339
+ "language_model.model.layers.10.linear_attn.in_proj_qkv": {
340
+ "bits": 2,
341
+ "group_size": 64,
342
+ "rank": 16
343
+ },
344
+ "language_model.model.layers.11.mlp.up_proj": {
345
+ "bits": 2,
346
+ "group_size": 64,
347
+ "rank": 16
348
+ },
349
+ "language_model.model.layers.11.mlp.down_proj": {
350
+ "bits": 2,
351
+ "group_size": 64,
352
+ "rank": 16
353
+ },
354
+ "language_model.model.layers.11.mlp.gate_proj": {
355
+ "bits": 2,
356
+ "group_size": 64,
357
+ "rank": 16
358
+ },
359
+ "language_model.model.layers.11.self_attn.o_proj": {
360
+ "bits": 4,
361
+ "group_size": 64,
362
+ "rank": 16
363
+ },
364
+ "language_model.model.layers.11.self_attn.v_proj": {
365
+ "bits": 4,
366
+ "group_size": 64,
367
+ "rank": 16
368
+ },
369
+ "language_model.model.layers.11.self_attn.k_proj": {
370
+ "bits": 4,
371
+ "group_size": 64,
372
+ "rank": 16
373
+ },
374
+ "language_model.model.layers.11.self_attn.q_proj": {
375
+ "bits": 4,
376
+ "group_size": 64,
377
+ "rank": 16
378
+ },
379
+ "language_model.model.layers.12.mlp.up_proj": {
380
+ "bits": 2,
381
+ "group_size": 64,
382
+ "rank": 16
383
+ },
384
+ "language_model.model.layers.12.mlp.down_proj": {
385
+ "bits": 2,
386
+ "group_size": 64,
387
+ "rank": 16
388
+ },
389
+ "language_model.model.layers.12.mlp.gate_proj": {
390
+ "bits": 2,
391
+ "group_size": 64,
392
+ "rank": 16
393
+ },
394
+ "language_model.model.layers.12.linear_attn.out_proj": {
395
+ "bits": 4,
396
+ "group_size": 64,
397
+ "rank": 16
398
+ },
399
+ "language_model.model.layers.12.linear_attn.in_proj_z": {
400
+ "bits": 2,
401
+ "group_size": 64,
402
+ "rank": 16
403
+ },
404
+ "language_model.model.layers.12.linear_attn.in_proj_qkv": {
405
+ "bits": 2,
406
+ "group_size": 64,
407
+ "rank": 16
408
+ },
409
+ "language_model.model.layers.13.mlp.up_proj": {
410
+ "bits": 2,
411
+ "group_size": 64,
412
+ "rank": 16
413
+ },
414
+ "language_model.model.layers.13.mlp.down_proj": {
415
+ "bits": 2,
416
+ "group_size": 64,
417
+ "rank": 16
418
+ },
419
+ "language_model.model.layers.13.mlp.gate_proj": {
420
+ "bits": 2,
421
+ "group_size": 64,
422
+ "rank": 16
423
+ },
424
+ "language_model.model.layers.13.linear_attn.out_proj": {
425
+ "bits": 4,
426
+ "group_size": 64,
427
+ "rank": 16
428
+ },
429
+ "language_model.model.layers.13.linear_attn.in_proj_z": {
430
+ "bits": 2,
431
+ "group_size": 64,
432
+ "rank": 16
433
+ },
434
+ "language_model.model.layers.13.linear_attn.in_proj_qkv": {
435
+ "bits": 2,
436
+ "group_size": 64,
437
+ "rank": 16
438
+ },
439
+ "language_model.model.layers.14.mlp.up_proj": {
440
+ "bits": 2,
441
+ "group_size": 64,
442
+ "rank": 16
443
+ },
444
+ "language_model.model.layers.14.mlp.down_proj": {
445
+ "bits": 2,
446
+ "group_size": 64,
447
+ "rank": 16
448
+ },
449
+ "language_model.model.layers.14.mlp.gate_proj": {
450
+ "bits": 2,
451
+ "group_size": 64,
452
+ "rank": 16
453
+ },
454
+ "language_model.model.layers.14.linear_attn.out_proj": {
455
+ "bits": 4,
456
+ "group_size": 64,
457
+ "rank": 16
458
+ },
459
+ "language_model.model.layers.14.linear_attn.in_proj_z": {
460
+ "bits": 2,
461
+ "group_size": 64,
462
+ "rank": 16
463
+ },
464
+ "language_model.model.layers.14.linear_attn.in_proj_qkv": {
465
+ "bits": 2,
466
+ "group_size": 64,
467
+ "rank": 16
468
+ },
469
+ "language_model.model.layers.15.mlp.up_proj": {
470
+ "bits": 2,
471
+ "group_size": 64,
472
+ "rank": 16
473
+ },
474
+ "language_model.model.layers.15.mlp.down_proj": {
475
+ "bits": 2,
476
+ "group_size": 64,
477
+ "rank": 16
478
+ },
479
+ "language_model.model.layers.15.mlp.gate_proj": {
480
+ "bits": 2,
481
+ "group_size": 64,
482
+ "rank": 16
483
+ },
484
+ "language_model.model.layers.15.self_attn.o_proj": {
485
+ "bits": 4,
486
+ "group_size": 64,
487
+ "rank": 16
488
+ },
489
+ "language_model.model.layers.15.self_attn.v_proj": {
490
+ "bits": 4,
491
+ "group_size": 64,
492
+ "rank": 16
493
+ },
494
+ "language_model.model.layers.15.self_attn.k_proj": {
495
+ "bits": 4,
496
+ "group_size": 64,
497
+ "rank": 16
498
+ },
499
+ "language_model.model.layers.15.self_attn.q_proj": {
500
+ "bits": 4,
501
+ "group_size": 64,
502
+ "rank": 16
503
+ },
504
+ "language_model.model.layers.16.mlp.up_proj": {
505
+ "bits": 2,
506
+ "group_size": 64,
507
+ "rank": 16
508
+ },
509
+ "language_model.model.layers.16.mlp.down_proj": {
510
+ "bits": 2,
511
+ "group_size": 64,
512
+ "rank": 16
513
+ },
514
+ "language_model.model.layers.16.mlp.gate_proj": {
515
+ "bits": 2,
516
+ "group_size": 64,
517
+ "rank": 16
518
+ },
519
+ "language_model.model.layers.16.linear_attn.out_proj": {
520
+ "bits": 4,
521
+ "group_size": 64,
522
+ "rank": 16
523
+ },
524
+ "language_model.model.layers.16.linear_attn.in_proj_z": {
525
+ "bits": 2,
526
+ "group_size": 64,
527
+ "rank": 16
528
+ },
529
+ "language_model.model.layers.16.linear_attn.in_proj_qkv": {
530
+ "bits": 2,
531
+ "group_size": 64,
532
+ "rank": 16
533
+ },
534
+ "language_model.model.layers.17.mlp.up_proj": {
535
+ "bits": 2,
536
+ "group_size": 64,
537
+ "rank": 16
538
+ },
539
+ "language_model.model.layers.17.mlp.down_proj": {
540
+ "bits": 2,
541
+ "group_size": 64,
542
+ "rank": 16
543
+ },
544
+ "language_model.model.layers.17.mlp.gate_proj": {
545
+ "bits": 2,
546
+ "group_size": 64,
547
+ "rank": 16
548
+ },
549
+ "language_model.model.layers.17.linear_attn.out_proj": {
550
+ "bits": 4,
551
+ "group_size": 64,
552
+ "rank": 16
553
+ },
554
+ "language_model.model.layers.17.linear_attn.in_proj_z": {
555
+ "bits": 2,
556
+ "group_size": 64,
557
+ "rank": 16
558
+ },
559
+ "language_model.model.layers.17.linear_attn.in_proj_qkv": {
560
+ "bits": 2,
561
+ "group_size": 64,
562
+ "rank": 16
563
+ },
564
+ "language_model.model.layers.18.mlp.up_proj": {
565
+ "bits": 2,
566
+ "group_size": 64,
567
+ "rank": 16
568
+ },
569
+ "language_model.model.layers.18.mlp.down_proj": {
570
+ "bits": 2,
571
+ "group_size": 64,
572
+ "rank": 16
573
+ },
574
+ "language_model.model.layers.18.mlp.gate_proj": {
575
+ "bits": 2,
576
+ "group_size": 64,
577
+ "rank": 16
578
+ },
579
+ "language_model.model.layers.18.linear_attn.out_proj": {
580
+ "bits": 4,
581
+ "group_size": 64,
582
+ "rank": 16
583
+ },
584
+ "language_model.model.layers.18.linear_attn.in_proj_z": {
585
+ "bits": 2,
586
+ "group_size": 64,
587
+ "rank": 16
588
+ },
589
+ "language_model.model.layers.18.linear_attn.in_proj_qkv": {
590
+ "bits": 2,
591
+ "group_size": 64,
592
+ "rank": 16
593
+ },
594
+ "language_model.model.layers.19.mlp.up_proj": {
595
+ "bits": 2,
596
+ "group_size": 64,
597
+ "rank": 16
598
+ },
599
+ "language_model.model.layers.19.mlp.down_proj": {
600
+ "bits": 2,
601
+ "group_size": 64,
602
+ "rank": 16
603
+ },
604
+ "language_model.model.layers.19.mlp.gate_proj": {
605
+ "bits": 2,
606
+ "group_size": 64,
607
+ "rank": 16
608
+ },
609
+ "language_model.model.layers.19.self_attn.o_proj": {
610
+ "bits": 4,
611
+ "group_size": 64,
612
+ "rank": 16
613
+ },
614
+ "language_model.model.layers.19.self_attn.v_proj": {
615
+ "bits": 4,
616
+ "group_size": 64,
617
+ "rank": 16
618
+ },
619
+ "language_model.model.layers.19.self_attn.k_proj": {
620
+ "bits": 4,
621
+ "group_size": 64,
622
+ "rank": 16
623
+ },
624
+ "language_model.model.layers.19.self_attn.q_proj": {
625
+ "bits": 4,
626
+ "group_size": 64,
627
+ "rank": 16
628
+ },
629
+ "language_model.model.layers.20.mlp.up_proj": {
630
+ "bits": 2,
631
+ "group_size": 64,
632
+ "rank": 16
633
+ },
634
+ "language_model.model.layers.20.mlp.down_proj": {
635
+ "bits": 2,
636
+ "group_size": 64,
637
+ "rank": 16
638
+ },
639
+ "language_model.model.layers.20.mlp.gate_proj": {
640
+ "bits": 2,
641
+ "group_size": 64,
642
+ "rank": 16
643
+ },
644
+ "language_model.model.layers.20.linear_attn.out_proj": {
645
+ "bits": 4,
646
+ "group_size": 64,
647
+ "rank": 16
648
+ },
649
+ "language_model.model.layers.20.linear_attn.in_proj_z": {
650
+ "bits": 2,
651
+ "group_size": 64,
652
+ "rank": 16
653
+ },
654
+ "language_model.model.layers.20.linear_attn.in_proj_qkv": {
655
+ "bits": 2,
656
+ "group_size": 64,
657
+ "rank": 16
658
+ },
659
+ "language_model.model.layers.21.mlp.up_proj": {
660
+ "bits": 2,
661
+ "group_size": 64,
662
+ "rank": 16
663
+ },
664
+ "language_model.model.layers.21.mlp.down_proj": {
665
+ "bits": 2,
666
+ "group_size": 64,
667
+ "rank": 16
668
+ },
669
+ "language_model.model.layers.21.mlp.gate_proj": {
670
+ "bits": 2,
671
+ "group_size": 64,
672
+ "rank": 16
673
+ },
674
+ "language_model.model.layers.21.linear_attn.out_proj": {
675
+ "bits": 4,
676
+ "group_size": 64,
677
+ "rank": 16
678
+ },
679
+ "language_model.model.layers.21.linear_attn.in_proj_z": {
680
+ "bits": 2,
681
+ "group_size": 64,
682
+ "rank": 16
683
+ },
684
+ "language_model.model.layers.21.linear_attn.in_proj_qkv": {
685
+ "bits": 2,
686
+ "group_size": 64,
687
+ "rank": 16
688
+ },
689
+ "language_model.model.layers.22.mlp.up_proj": {
690
+ "bits": 2,
691
+ "group_size": 64,
692
+ "rank": 16
693
+ },
694
+ "language_model.model.layers.22.mlp.down_proj": {
695
+ "bits": 2,
696
+ "group_size": 64,
697
+ "rank": 16
698
+ },
699
+ "language_model.model.layers.22.mlp.gate_proj": {
700
+ "bits": 2,
701
+ "group_size": 64,
702
+ "rank": 16
703
+ },
704
+ "language_model.model.layers.22.linear_attn.out_proj": {
705
+ "bits": 4,
706
+ "group_size": 64,
707
+ "rank": 16
708
+ },
709
+ "language_model.model.layers.22.linear_attn.in_proj_z": {
710
+ "bits": 2,
711
+ "group_size": 64,
712
+ "rank": 16
713
+ },
714
+ "language_model.model.layers.22.linear_attn.in_proj_qkv": {
715
+ "bits": 2,
716
+ "group_size": 64,
717
+ "rank": 16
718
+ },
719
+ "language_model.model.layers.23.mlp.up_proj": {
720
+ "bits": 2,
721
+ "group_size": 64,
722
+ "rank": 16
723
+ },
724
+ "language_model.model.layers.23.mlp.down_proj": {
725
+ "bits": 2,
726
+ "group_size": 64,
727
+ "rank": 16
728
+ },
729
+ "language_model.model.layers.23.mlp.gate_proj": {
730
+ "bits": 2,
731
+ "group_size": 64,
732
+ "rank": 16
733
+ },
734
+ "language_model.model.layers.23.self_attn.o_proj": {
735
+ "bits": 4,
736
+ "group_size": 64,
737
+ "rank": 16
738
+ },
739
+ "language_model.model.layers.23.self_attn.v_proj": {
740
+ "bits": 4,
741
+ "group_size": 64,
742
+ "rank": 16
743
+ },
744
+ "language_model.model.layers.23.self_attn.k_proj": {
745
+ "bits": 4,
746
+ "group_size": 64,
747
+ "rank": 16
748
+ },
749
+ "language_model.model.layers.23.self_attn.q_proj": {
750
+ "bits": 4,
751
+ "group_size": 64,
752
+ "rank": 16
753
+ },
754
+ "language_model.model.layers.24.mlp.up_proj": {
755
+ "bits": 2,
756
+ "group_size": 64,
757
+ "rank": 16
758
+ },
759
+ "language_model.model.layers.24.mlp.down_proj": {
760
+ "bits": 2,
761
+ "group_size": 64,
762
+ "rank": 16
763
+ },
764
+ "language_model.model.layers.24.mlp.gate_proj": {
765
+ "bits": 2,
766
+ "group_size": 64,
767
+ "rank": 16
768
+ },
769
+ "language_model.model.layers.24.linear_attn.out_proj": {
770
+ "bits": 4,
771
+ "group_size": 64,
772
+ "rank": 16
773
+ },
774
+ "language_model.model.layers.24.linear_attn.in_proj_z": {
775
+ "bits": 2,
776
+ "group_size": 64,
777
+ "rank": 16
778
+ },
779
+ "language_model.model.layers.24.linear_attn.in_proj_qkv": {
780
+ "bits": 2,
781
+ "group_size": 64,
782
+ "rank": 16
783
+ },
784
+ "language_model.model.layers.25.mlp.up_proj": {
785
+ "bits": 2,
786
+ "group_size": 64,
787
+ "rank": 16
788
+ },
789
+ "language_model.model.layers.25.mlp.down_proj": {
790
+ "bits": 2,
791
+ "group_size": 64,
792
+ "rank": 16
793
+ },
794
+ "language_model.model.layers.25.mlp.gate_proj": {
795
+ "bits": 2,
796
+ "group_size": 64,
797
+ "rank": 16
798
+ },
799
+ "language_model.model.layers.25.linear_attn.out_proj": {
800
+ "bits": 4,
801
+ "group_size": 64,
802
+ "rank": 16
803
+ },
804
+ "language_model.model.layers.25.linear_attn.in_proj_z": {
805
+ "bits": 2,
806
+ "group_size": 64,
807
+ "rank": 16
808
+ },
809
+ "language_model.model.layers.25.linear_attn.in_proj_qkv": {
810
+ "bits": 2,
811
+ "group_size": 64,
812
+ "rank": 16
813
+ },
814
+ "language_model.model.layers.26.mlp.up_proj": {
815
+ "bits": 2,
816
+ "group_size": 64,
817
+ "rank": 16
818
+ },
819
+ "language_model.model.layers.26.mlp.down_proj": {
820
+ "bits": 2,
821
+ "group_size": 64,
822
+ "rank": 16
823
+ },
824
+ "language_model.model.layers.26.mlp.gate_proj": {
825
+ "bits": 2,
826
+ "group_size": 64,
827
+ "rank": 16
828
+ },
829
+ "language_model.model.layers.26.linear_attn.out_proj": {
830
+ "bits": 4,
831
+ "group_size": 64,
832
+ "rank": 16
833
+ },
834
+ "language_model.model.layers.26.linear_attn.in_proj_z": {
835
+ "bits": 2,
836
+ "group_size": 64,
837
+ "rank": 16
838
+ },
839
+ "language_model.model.layers.26.linear_attn.in_proj_qkv": {
840
+ "bits": 2,
841
+ "group_size": 64,
842
+ "rank": 16
843
+ },
844
+ "language_model.model.layers.27.mlp.up_proj": {
845
+ "bits": 2,
846
+ "group_size": 64,
847
+ "rank": 16
848
+ },
849
+ "language_model.model.layers.27.mlp.down_proj": {
850
+ "bits": 2,
851
+ "group_size": 64,
852
+ "rank": 16
853
+ },
854
+ "language_model.model.layers.27.mlp.gate_proj": {
855
+ "bits": 2,
856
+ "group_size": 64,
857
+ "rank": 16
858
+ },
859
+ "language_model.model.layers.27.self_attn.o_proj": {
860
+ "bits": 4,
861
+ "group_size": 64,
862
+ "rank": 16
863
+ },
864
+ "language_model.model.layers.27.self_attn.v_proj": {
865
+ "bits": 4,
866
+ "group_size": 64,
867
+ "rank": 16
868
+ },
869
+ "language_model.model.layers.27.self_attn.k_proj": {
870
+ "bits": 4,
871
+ "group_size": 64,
872
+ "rank": 16
873
+ },
874
+ "language_model.model.layers.27.self_attn.q_proj": {
875
+ "bits": 4,
876
+ "group_size": 64,
877
+ "rank": 16
878
+ },
879
+ "language_model.model.layers.28.mlp.up_proj": {
880
+ "bits": 2,
881
+ "group_size": 64,
882
+ "rank": 16
883
+ },
884
+ "language_model.model.layers.28.mlp.down_proj": {
885
+ "bits": 2,
886
+ "group_size": 64,
887
+ "rank": 16
888
+ },
889
+ "language_model.model.layers.28.mlp.gate_proj": {
890
+ "bits": 2,
891
+ "group_size": 64,
892
+ "rank": 16
893
+ },
894
+ "language_model.model.layers.28.linear_attn.out_proj": {
895
+ "bits": 4,
896
+ "group_size": 64,
897
+ "rank": 16
898
+ },
899
+ "language_model.model.layers.28.linear_attn.in_proj_z": {
900
+ "bits": 2,
901
+ "group_size": 64,
902
+ "rank": 16
903
+ },
904
+ "language_model.model.layers.28.linear_attn.in_proj_qkv": {
905
+ "bits": 2,
906
+ "group_size": 64,
907
+ "rank": 16
908
+ },
909
+ "language_model.model.layers.29.mlp.up_proj": {
910
+ "bits": 2,
911
+ "group_size": 64,
912
+ "rank": 16
913
+ },
914
+ "language_model.model.layers.29.mlp.down_proj": {
915
+ "bits": 2,
916
+ "group_size": 64,
917
+ "rank": 16
918
+ },
919
+ "language_model.model.layers.29.mlp.gate_proj": {
920
+ "bits": 2,
921
+ "group_size": 64,
922
+ "rank": 16
923
+ },
924
+ "language_model.model.layers.29.linear_attn.out_proj": {
925
+ "bits": 4,
926
+ "group_size": 64,
927
+ "rank": 16
928
+ },
929
+ "language_model.model.layers.29.linear_attn.in_proj_z": {
930
+ "bits": 2,
931
+ "group_size": 64,
932
+ "rank": 16
933
+ },
934
+ "language_model.model.layers.29.linear_attn.in_proj_qkv": {
935
+ "bits": 2,
936
+ "group_size": 64,
937
+ "rank": 16
938
+ },
939
+ "language_model.model.layers.30.mlp.up_proj": {
940
+ "bits": 2,
941
+ "group_size": 64,
942
+ "rank": 16
943
+ },
944
+ "language_model.model.layers.30.mlp.down_proj": {
945
+ "bits": 2,
946
+ "group_size": 64,
947
+ "rank": 16
948
+ },
949
+ "language_model.model.layers.30.mlp.gate_proj": {
950
+ "bits": 2,
951
+ "group_size": 64,
952
+ "rank": 16
953
+ },
954
+ "language_model.model.layers.30.linear_attn.out_proj": {
955
+ "bits": 4,
956
+ "group_size": 64,
957
+ "rank": 16
958
+ },
959
+ "language_model.model.layers.30.linear_attn.in_proj_z": {
960
+ "bits": 2,
961
+ "group_size": 64,
962
+ "rank": 16
963
+ },
964
+ "language_model.model.layers.30.linear_attn.in_proj_qkv": {
965
+ "bits": 2,
966
+ "group_size": 64,
967
+ "rank": 16
968
+ },
969
+ "language_model.model.layers.31.mlp.up_proj": {
970
+ "bits": 2,
971
+ "group_size": 64,
972
+ "rank": 16
973
+ },
974
+ "language_model.model.layers.31.mlp.down_proj": {
975
+ "bits": 2,
976
+ "group_size": 64,
977
+ "rank": 16
978
+ },
979
+ "language_model.model.layers.31.mlp.gate_proj": {
980
+ "bits": 2,
981
+ "group_size": 64,
982
+ "rank": 16
983
+ },
984
+ "language_model.model.layers.31.self_attn.o_proj": {
985
+ "bits": 4,
986
+ "group_size": 64,
987
+ "rank": 16
988
+ },
989
+ "language_model.model.layers.31.self_attn.v_proj": {
990
+ "bits": 4,
991
+ "group_size": 64,
992
+ "rank": 16
993
+ },
994
+ "language_model.model.layers.31.self_attn.k_proj": {
995
+ "bits": 4,
996
+ "group_size": 64,
997
+ "rank": 16
998
+ },
999
+ "language_model.model.layers.31.self_attn.q_proj": {
1000
+ "bits": 4,
1001
+ "group_size": 64,
1002
+ "rank": 16
1003
+ }
1004
+ },
1005
+ "embedding": {
1006
+ "bits": 4,
1007
+ "group_size": 64
1008
+ },
1009
+ "source": {
1010
+ "format": "kev-packed-mlx-v1",
1011
+ "source_packed_sha256": "af8703c3c8f848c18f49a40e91978363fae1e0e3b6c4678a9436e3ebc5da73e0",
1012
+ "source_cloq_sha256": "52b8699b25a819542d8073b0de52c79bfb77e46b3bc969e6f553b6a047e0afcc",
1013
+ "arithmetic": "native affine quantized matmul plus separate fp32 low-rank compensation; not merged-bf16 parity"
1014
+ },
1015
+ "model_sha256": "c4899eaa9fb206e77260f832d00188c9b90750724422ea0eea43ed69faeff43d"
1016
+ }
release-manifest.json ADDED
@@ -0,0 +1,252 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "repo_id": "Jakevin/kev-4b-ternary-mlx",
3
+ "tag": "v0.1.0-experimental",
4
+ "date": "2026-10-06",
5
+ "format": "kev-packed-mlx-v1",
6
+ "source_kev_revision": "6cfce5c2fa4b4bd64026336ab649c5ca78857d52",
7
+ "files": {
8
+ "LICENSE": {
9
+ "bytes": 11343,
10
+ "sha256": "6b08bb37982c233aa12bcbdf19106da12f3f4fcf800773ea42e28ebeddd34fb8"
11
+ },
12
+ "NOTICE": {
13
+ "bytes": 528,
14
+ "sha256": "9021541d1f4af146262ef38449dd517c13d059404132e3f7e017ab550cef62bc"
15
+ },
16
+ "README.md": {
17
+ "bytes": 8520,
18
+ "sha256": "c11cf88162f1739d10f5e68f1a1d2bdd05d794381ebb99dbaa65a06f0f0c4dea"
19
+ },
20
+ "added_tokens.json": {
21
+ "bytes": 707,
22
+ "sha256": "c0284b582e14987fbd3d5a2cb2bd139084371ed9acbae488829a1c900833c680"
23
+ },
24
+ "config.json": {
25
+ "bytes": 2692,
26
+ "sha256": "fc1557ea241bfc2c8abd87c1213729a44562fe5da9edf0fd5dc9f95b6cdbdaed"
27
+ },
28
+ "evaluation/development-summary.json": {
29
+ "bytes": 2475,
30
+ "sha256": "d17ad38950313c7513bd67a73cb2b15358c0a1f42dbec2db02c6c9388a392f5f"
31
+ },
32
+ "evaluation/memory-long.json": {
33
+ "bytes": 895,
34
+ "sha256": "b551f722360e3420a30aed1fada9d44557bef26730938df1a2a73b961a54354d"
35
+ },
36
+ "evaluation/memory-short.json": {
37
+ "bytes": 894,
38
+ "sha256": "d3111a799b8760323eba78cae3bf00fd61c61d24073a232cff33a9c7baf3d41d"
39
+ },
40
+ "evaluation/native-comparison.json": {
41
+ "bytes": 275,
42
+ "sha256": "a1355c9afa8b8f7f986c355941233164457d05503810d49fd698bd224adf8fc5"
43
+ },
44
+ "head.pt": {
45
+ "bytes": 5247807,
46
+ "sha256": "6f6740d94a49187770b68c0f058d85bedc0cf7efc20dff7ca9a418c8b461dc0e"
47
+ },
48
+ "merges.txt": {
49
+ "bytes": 1671853,
50
+ "sha256": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5"
51
+ },
52
+ "model.safetensors": {
53
+ "bytes": 1678909327,
54
+ "sha256": "c4899eaa9fb206e77260f832d00188c9b90750724422ea0eea43ed69faeff43d"
55
+ },
56
+ "packed.json": {
57
+ "bytes": 24929,
58
+ "sha256": "1ae5d2d98e967e49ebdf29842da0847faf87fa03ace9370f8e05715ef583dd29"
59
+ },
60
+ "runtime/LICENSE": {
61
+ "bytes": 11343,
62
+ "sha256": "6b08bb37982c233aa12bcbdf19106da12f3f4fcf800773ea42e28ebeddd34fb8"
63
+ },
64
+ "runtime/README.md": {
65
+ "bytes": 42304,
66
+ "sha256": "19c9d7d30cd723b9f112f1f68d0b415a28c5ee616dbc21b358fc29bd07a66311"
67
+ },
68
+ "runtime/docs/packed-mlx.md": {
69
+ "bytes": 6233,
70
+ "sha256": "c70c572479cb7e8e8094e7cdd5982ade28084c9468e38190ef0d13ffa675631b"
71
+ },
72
+ "runtime/kev/__init__.py": {
73
+ "bytes": 0,
74
+ "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
75
+ },
76
+ "runtime/kev/anchors.py": {
77
+ "bytes": 3930,
78
+ "sha256": "6963eafdb276db5a6c94d939eae675c761555553b8d0f963ffd803448beaabbb"
79
+ },
80
+ "runtime/kev/api.py": {
81
+ "bytes": 7526,
82
+ "sha256": "a6f53af5354f8ea17915abc78360fdf291a5108ce39a53a2d9135d48c603ae58"
83
+ },
84
+ "runtime/kev/autoresearch.py": {
85
+ "bytes": 11866,
86
+ "sha256": "7c62762b99c1b081bef06aa1ce83f8dfccce8ac7e9d2c5c0ab4d058055c7eab1"
87
+ },
88
+ "runtime/kev/benchmark.py": {
89
+ "bytes": 14203,
90
+ "sha256": "82e6a935472f170da48ec37eee8986c519595c4943f6f3040ec67b5da3afb278"
91
+ },
92
+ "runtime/kev/budget.py": {
93
+ "bytes": 8457,
94
+ "sha256": "246b5ca6450567a61b881308c9482cf1af6ecba599ca161af1febb23d6dc28a7"
95
+ },
96
+ "runtime/kev/calibrate.py": {
97
+ "bytes": 4846,
98
+ "sha256": "435ac98696fbd329c3a07f0c12633ba3c09a846f0e9c92e17540bbf8ba2986f8"
99
+ },
100
+ "runtime/kev/checkpoint.py": {
101
+ "bytes": 24924,
102
+ "sha256": "efe96cde9ef989c4fea3b6a0b631338bb1870bfe4b0f1eb99d71bb0761af34b7"
103
+ },
104
+ "runtime/kev/compare.py": {
105
+ "bytes": 3198,
106
+ "sha256": "3606dbf98bb305d66420158498cb837304d2edb8dab01cdaf2ec42964bef3435"
107
+ },
108
+ "runtime/kev/composition.py": {
109
+ "bytes": 15847,
110
+ "sha256": "f335ed17e18e0a544893db5e22b9a059e6ce1b2e14dbb863ac7d7bbf8f3e0536"
111
+ },
112
+ "runtime/kev/contrastive.py": {
113
+ "bytes": 20624,
114
+ "sha256": "cbb979aa5d40265ad0e64695f94b281d91751fa405811ddfa8212fede111edcf"
115
+ },
116
+ "runtime/kev/cuda_graphs.py": {
117
+ "bytes": 25356,
118
+ "sha256": "5ce84f0e4e5ad2177314ed3ecc9564d45d14ad9eb422e0fb079c9e1cbbb011b6"
119
+ },
120
+ "runtime/kev/data.py": {
121
+ "bytes": 25861,
122
+ "sha256": "e77a2fffeb8ee7e05b118893be8e38b8358aef97c3ca2e143cdce208201b92f9"
123
+ },
124
+ "runtime/kev/device.py": {
125
+ "bytes": 1164,
126
+ "sha256": "bda81dff5be33c59eee8de76e534976cc137a2661da6afd2a67abb85d8dad4a2"
127
+ },
128
+ "runtime/kev/evaluate.py": {
129
+ "bytes": 14408,
130
+ "sha256": "b2a6fbad7a1494728ea8d9ed0ad54ae23ff2ee1d543444fd0c3c110da4ab07cc"
131
+ },
132
+ "runtime/kev/experiment.py": {
133
+ "bytes": 34376,
134
+ "sha256": "3fd13a86f111dfd3292af9448df0ee28bd9596e3a77c1391f318cd654be14e8a"
135
+ },
136
+ "runtime/kev/full_ft.py": {
137
+ "bytes": 28171,
138
+ "sha256": "c3abef39cd9d486b3c075a37286afbcd1aa42a04c78a1b7b9b2e77b5a3963929"
139
+ },
140
+ "runtime/kev/fused_qwen35.py": {
141
+ "bytes": 10294,
142
+ "sha256": "5d2e587808822ce539ff52f63024357630b67a65f88cd734b82fd488655aebed"
143
+ },
144
+ "runtime/kev/jev.py": {
145
+ "bytes": 3275,
146
+ "sha256": "5488645c4fc2a53ca430e556ebfb744736e1a156bd121bd10ba0a9935f423e1c"
147
+ },
148
+ "runtime/kev/metrics.py": {
149
+ "bytes": 24401,
150
+ "sha256": "167a586334cfcced4adcc197758535c3bcda832b855c151095ef2d1968cf8983"
151
+ },
152
+ "runtime/kev/mirror.py": {
153
+ "bytes": 4982,
154
+ "sha256": "ca2284c670ce2cf1db57231968b9e6cbc01955f8f3c9a3207fd82e73d7b476b6"
155
+ },
156
+ "runtime/kev/mlx_model.py": {
157
+ "bytes": 15144,
158
+ "sha256": "8aef8aef93f50663395d73b5ae6cc964a01c4f4bf7f92a3843c1f32681cc9c09"
159
+ },
160
+ "runtime/kev/model.py": {
161
+ "bytes": 35153,
162
+ "sha256": "2ff3b2e67c048e3bc5a0d1a32aafe50804ed315a5564fce06bb4b95697c5271c"
163
+ },
164
+ "runtime/kev/packed_model.py": {
165
+ "bytes": 4892,
166
+ "sha256": "72cf553bd8d6f6e6d1e9f903eab4eba44d33d11571803752c0a89404f39e3499"
167
+ },
168
+ "runtime/kev/plot.py": {
169
+ "bytes": 3368,
170
+ "sha256": "0874bcff2885d8155a1cceade0de8a275a7163d3c3fd6aa0c294e8ccf5ff7e02"
171
+ },
172
+ "runtime/kev/predictors.py": {
173
+ "bytes": 21963,
174
+ "sha256": "29263ae3297e5cde6cf03da5c5aa3ed377c9681c6a5a639ea4aa259acec7575b"
175
+ },
176
+ "runtime/kev/publish.py": {
177
+ "bytes": 7512,
178
+ "sha256": "fcb490a374fc9ef330424ad4a75ce281ab08d2cb04dff2ad912c5f4dae6ab8a7"
179
+ },
180
+ "runtime/kev/rounds.py": {
181
+ "bytes": 96716,
182
+ "sha256": "2975bede7250812aecea36b17c0aefda6b389dc3a50e782e82e7566a7168c270"
183
+ },
184
+ "runtime/kev/serve.py": {
185
+ "bytes": 23057,
186
+ "sha256": "4f9885ce179dcb7de743f59be4b561758f47fac09373525e2994860b3345b162"
187
+ },
188
+ "runtime/kev/shared_prefix.py": {
189
+ "bytes": 9671,
190
+ "sha256": "bc6dedbad734cca227f3ec6746223b3b03363b0c045aa9f2493d886b1828e71e"
191
+ },
192
+ "runtime/kev/study_v3.py": {
193
+ "bytes": 14919,
194
+ "sha256": "7891150aa623e4479c7c789d4dc4662f186132b37b7bac1e1f072b9e08ee13e3"
195
+ },
196
+ "runtime/kev/suite.py": {
197
+ "bytes": 26388,
198
+ "sha256": "8e7d257a10b4704d7bf4f8ec2c74dddf26aa5a2a610ea78610360e6baf43eff5"
199
+ },
200
+ "runtime/kev/train.py": {
201
+ "bytes": 53543,
202
+ "sha256": "de0b0971defd79c053d56f8e8b17b0501b2e264e93a66a8d5493f1b1ba839416"
203
+ },
204
+ "runtime/kev/transfer_v9.py": {
205
+ "bytes": 10969,
206
+ "sha256": "0902409742151250a28af1fd8f42258b70fcc161f3deed3ee35fe3f473f3c763"
207
+ },
208
+ "runtime/pyproject.toml": {
209
+ "bytes": 1721,
210
+ "sha256": "7c17fbe9efcadda3eb488b59adf1b6db5938cd33af4d3bd389ce4359b73fcc7d"
211
+ },
212
+ "runtime/scripts/build_packed_model.py": {
213
+ "bytes": 7130,
214
+ "sha256": "1f4aba3f4c754753c3c654b934cd410acc7fb9c0f27fa35366e9d4a1899e5d5c"
215
+ },
216
+ "runtime/scripts/mlx_long_states.py": {
217
+ "bytes": 29651,
218
+ "sha256": "ec1ae5ff040da1060143dcbcd551bdd5982e583951af58677cdb0b6624ed005d"
219
+ },
220
+ "runtime/scripts/profile_packed_model.py": {
221
+ "bytes": 3543,
222
+ "sha256": "14b8e346771f978fb61b2d167cbb5f56011600f3e59ae4a16b84518fd99d2a18"
223
+ },
224
+ "runtime/tests/test_packed_model.py": {
225
+ "bytes": 7193,
226
+ "sha256": "9ff620541b964886f1e13117ffb4a6e143b3c843a144fd2fcc1a087baa3c82de"
227
+ },
228
+ "special_tokens_map.json": {
229
+ "bytes": 616,
230
+ "sha256": "6676f091c8bc4d1b50146427cfde92073402866b87b6e39223227931b70083e9"
231
+ },
232
+ "tokenizer.json": {
233
+ "bytes": 19989325,
234
+ "sha256": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523"
235
+ },
236
+ "tokenizer_config.json": {
237
+ "bytes": 1128,
238
+ "sha256": "8671bed7c852ce9e661be94f179a7b4ffd091c2a65aea0363e5501c20318ee45"
239
+ },
240
+ "vocab.json": {
241
+ "bytes": 2776833,
242
+ "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"
243
+ }
244
+ },
245
+ "excluded": [
246
+ "original backbone and PEFT weights",
247
+ "NAER/Chinese Linguipedia sources, derivatives, labels and predictions",
248
+ "training and validation JSONL",
249
+ "per-example evaluation prediction dumps"
250
+ ],
251
+ "note": "Runtime code included; model SHA verified against packed.json. Manifest excludes its own hash."
252
+ }
runtime/LICENSE ADDED
@@ -0,0 +1,202 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ Apache License
3
+ Version 2.0, January 2004
4
+ http://www.apache.org/licenses/
5
+
6
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
7
+
8
+ 1. Definitions.
9
+
10
+ "License" shall mean the terms and conditions for use, reproduction,
11
+ and distribution as defined by Sections 1 through 9 of this document.
12
+
13
+ "Licensor" shall mean the copyright owner or entity authorized by
14
+ the copyright owner that is granting the License.
15
+
16
+ "Legal Entity" shall mean the union of the acting entity and all
17
+ other entities that control, are controlled by, or are under common
18
+ control with that entity. For the purposes of this definition,
19
+ "control" means (i) the power, direct or indirect, to cause the
20
+ direction or management of such entity, whether by contract or
21
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
22
+ outstanding shares, or (iii) beneficial ownership of such entity.
23
+
24
+ "You" (or "Your") shall mean an individual or Legal Entity
25
+ exercising permissions granted by this License.
26
+
27
+ "Source" form shall mean the preferred form for making modifications,
28
+ including but not limited to software source code, documentation
29
+ source, and configuration files.
30
+
31
+ "Object" form shall mean any form resulting from mechanical
32
+ transformation or translation of a Source form, including but
33
+ not limited to compiled object code, generated documentation,
34
+ and conversions to other media types.
35
+
36
+ "Work" shall mean the work of authorship, whether in Source or
37
+ Object form, made available under the License, as indicated by a
38
+ copyright notice that is included in or attached to the work
39
+ (an example is provided in the Appendix below).
40
+
41
+ "Derivative Works" shall mean any work, whether in Source or Object
42
+ form, that is based on (or derived from) the Work and for which the
43
+ editorial revisions, annotations, elaborations, or other modifications
44
+ represent, as a whole, an original work of authorship. For the purposes
45
+ of this License, Derivative Works shall not include works that remain
46
+ separable from, or merely link (or bind by name) to the interfaces of,
47
+ the Work and Derivative Works thereof.
48
+
49
+ "Contribution" shall mean any work of authorship, including
50
+ the original version of the Work and any modifications or additions
51
+ to that Work or Derivative Works thereof, that is intentionally
52
+ submitted to Licensor for inclusion in the Work by the copyright owner
53
+ or by an individual or Legal Entity authorized to submit on behalf of
54
+ the copyright owner. For the purposes of this definition, "submitted"
55
+ means any form of electronic, verbal, or written communication sent
56
+ to the Licensor or its representatives, including but not limited to
57
+ communication on electronic mailing lists, source code control systems,
58
+ and issue tracking systems that are managed by, or on behalf of, the
59
+ Licensor for the purpose of discussing and improving the Work, but
60
+ excluding communication that is conspicuously marked or otherwise
61
+ designated in writing by the copyright owner as "Not a Contribution."
62
+
63
+ "Contributor" shall mean Licensor and any individual or Legal Entity
64
+ on behalf of whom a Contribution has been received by Licensor and
65
+ subsequently incorporated within the Work.
66
+
67
+ 2. Grant of Copyright License. Subject to the terms and conditions of
68
+ this License, each Contributor hereby grants to You a perpetual,
69
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
70
+ copyright license to reproduce, prepare Derivative Works of,
71
+ publicly display, publicly perform, sublicense, and distribute the
72
+ Work and such Derivative Works in Source or Object form.
73
+
74
+ 3. Grant of Patent License. Subject to the terms and conditions of
75
+ this License, each Contributor hereby grants to You a perpetual,
76
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
77
+ (except as stated in this section) patent license to make, have made,
78
+ use, offer to sell, sell, import, and otherwise transfer the Work,
79
+ where such license applies only to those patent claims licensable
80
+ by such Contributor that are necessarily infringed by their
81
+ Contribution(s) alone or by combination of their Contribution(s)
82
+ with the Work to which such Contribution(s) was submitted. If You
83
+ institute patent litigation against any entity (including a
84
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
85
+ or a Contribution incorporated within the Work constitutes direct
86
+ or contributory patent infringement, then any patent licenses
87
+ granted to You under this License for that Work shall terminate
88
+ as of the date such litigation is filed.
89
+
90
+ 4. Redistribution. You may reproduce and distribute copies of the
91
+ Work or Derivative Works thereof in any medium, with or without
92
+ modifications, and in Source or Object form, provided that You
93
+ meet the following conditions:
94
+
95
+ (a) You must give any other recipients of the Work or
96
+ Derivative Works a copy of this License; and
97
+
98
+ (b) You must cause any modified files to carry prominent notices
99
+ stating that You changed the files; and
100
+
101
+ (c) You must retain, in the Source form of any Derivative Works
102
+ that You distribute, all copyright, patent, trademark, and
103
+ attribution notices from the Source form of the Work,
104
+ excluding those notices that do not pertain to any part of
105
+ the Derivative Works; and
106
+
107
+ (d) If the Work includes a "NOTICE" text file as part of its
108
+ distribution, then any Derivative Works that You distribute must
109
+ include a readable copy of the attribution notices contained
110
+ within such NOTICE file, excluding those notices that do not
111
+ pertain to any part of the Derivative Works, in at least one
112
+ of the following places: within a NOTICE text file distributed
113
+ as part of the Derivative Works; within the Source form or
114
+ documentation, if provided along with the Derivative Works; or,
115
+ within a display generated by the Derivative Works, if and
116
+ wherever such third-party notices normally appear. The contents
117
+ of the NOTICE file are for informational purposes only and
118
+ do not modify the License. You may add Your own attribution
119
+ notices within Derivative Works that You distribute, alongside
120
+ or as an addendum to the NOTICE text from the Work, provided
121
+ that such additional attribution notices cannot be construed
122
+ as modifying the License.
123
+
124
+ You may add Your own copyright statement to Your modifications and
125
+ may provide additional or different license terms and conditions
126
+ for use, reproduction, or distribution of Your modifications, or
127
+ for any such Derivative Works as a whole, provided Your use,
128
+ reproduction, and distribution of the Work otherwise complies with
129
+ the conditions stated in this License.
130
+
131
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
132
+ any Contribution intentionally submitted for inclusion in the Work
133
+ by You to the Licensor shall be under the terms and conditions of
134
+ this License, without any additional terms or conditions.
135
+ Notwithstanding the above, nothing herein shall supersede or modify
136
+ the terms of any separate license agreement you may have executed
137
+ with Licensor regarding such Contributions.
138
+
139
+ 6. Trademarks. This License does not grant permission to use the trade
140
+ names, trademarks, service marks, or product names of the Licensor,
141
+ except as required for reasonable and customary use in describing the
142
+ origin of the Work and reproducing the content of the NOTICE file.
143
+
144
+ 7. Disclaimer of Warranty. Unless required by applicable law or
145
+ agreed to in writing, Licensor provides the Work (and each
146
+ Contributor provides its Contributions) on an "AS IS" BASIS,
147
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
148
+ implied, including, without limitation, any warranties or conditions
149
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
150
+ PARTICULAR PURPOSE. You are solely responsible for determining the
151
+ appropriateness of using or redistributing the Work and assume any
152
+ risks associated with Your exercise of permissions under this License.
153
+
154
+ 8. Limitation of Liability. In no event and under no legal theory,
155
+ whether in tort (including negligence), contract, or otherwise,
156
+ unless required by applicable law (such as deliberate and grossly
157
+ negligent acts) or agreed to in writing, shall any Contributor be
158
+ liable to You for damages, including any direct, indirect, special,
159
+ incidental, or consequential damages of any character arising as a
160
+ result of this License or out of the use or inability to use the
161
+ Work (including but not limited to damages for loss of goodwill,
162
+ work stoppage, computer failure or malfunction, or any and all
163
+ other commercial damages or losses), even if such Contributor
164
+ has been advised of the possibility of such damages.
165
+
166
+ 9. Accepting Warranty or Additional Liability. While redistributing
167
+ the Work or Derivative Works thereof, You may choose to offer,
168
+ and charge a fee for, acceptance of support, warranty, indemnity,
169
+ or other liability obligations and/or rights consistent with this
170
+ License. However, in accepting such obligations, You may act only
171
+ on Your own behalf and on Your sole responsibility, not on behalf
172
+ of any other Contributor, and only if You agree to indemnify,
173
+ defend, and hold each Contributor harmless for any liability
174
+ incurred by, or claims asserted against, such Contributor by reason
175
+ of your accepting any such warranty or additional liability.
176
+
177
+ END OF TERMS AND CONDITIONS
178
+
179
+ APPENDIX: How to apply the Apache License to your work.
180
+
181
+ To apply the Apache License to your work, attach the following
182
+ boilerplate notice, with the fields enclosed by brackets "[]"
183
+ replaced with your own identifying information. (Don't include
184
+ the brackets!) The text should be enclosed in the appropriate
185
+ comment syntax for the file format. We also recommend that a
186
+ file or class name and description of purpose be included on the
187
+ same "printed page" as the copyright notice for easier
188
+ identification within third-party archives.
189
+
190
+ Copyright 2026 Jared Palmer
191
+
192
+ Licensed under the Apache License, Version 2.0 (the "License");
193
+ you may not use this file except in compliance with the License.
194
+ You may obtain a copy of the License at
195
+
196
+ http://www.apache.org/licenses/LICENSE-2.0
197
+
198
+ Unless required by applicable law or agreed to in writing, software
199
+ distributed under the License is distributed on an "AS IS" BASIS,
200
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
201
+ See the License for the specific language governing permissions and
202
+ limitations under the License.
runtime/README.md ADDED
@@ -0,0 +1,466 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Kev
2
+
3
+ Small Jev-like decision models you can train and run yourself.
4
+
5
+ <p>
6
+ <a href="https://github.com/jaredpalmer/kev/actions/workflows/ci.yml"><img alt="CI" src="https://img.shields.io/github/actions/workflow/status/jaredpalmer/kev/ci.yml?style=for-the-badge&labelColor=000000" height="28"></a>
7
+ <a href="https://huggingface.co/collections/jaredpalmer/kev-6aad9d0ea49f2589665e07cd"><img alt="Weights: Kev-0.8B · 4B · 9B · 27B" src="https://img.shields.io/badge/WEIGHTS-0.8B%20%C2%B7%204B%20%C2%B7%209B%20%C2%B7%2027B-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
8
+ <a href="https://huggingface.co/spaces/jaredpalmer/kev"><img alt="Demo on Hugging Face Spaces" src="https://img.shields.io/badge/DEMO-HF%20Spaces-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
9
+ <a href="https://huggingface.co/datasets/jaredpalmer/kev-suites"><img alt="Frozen eval suites" src="https://img.shields.io/badge/EVAL%20SUITES-frozen-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
10
+ <a href="LICENSE"><img alt="License: Apache-2.0" src="https://img.shields.io/badge/license-Apache--2.0-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
11
+ </p>
12
+
13
+ Kev is a family of small decision models built on Qwen3.5 and Qwen3.8 and based on the architecture described in [Jev's Architecture Unmasked](https://archerhume.com/posts/jevs-architecture-unmasked). You can use the pretrained weights or train your own. The API matches TypeSafe's [System One](https://docs.typesafe.ai/api), so you can point their Python SDK at your local server.
14
+
15
+ ## Highlights
16
+
17
+ - Yes/no (`noul`), multiple-choice (`choice`) and rating (`score`) questions in one request. The questions share the text but can't read each other.
18
+ - Calibrated probabilities by default: each checkpoint ships with a fitted temperature.
19
+ - Drop-in for Jev: the TypeSafe Python SDK works against a Kev server unchanged.
20
+ - Four sizes, versioned together as Kev 1.0: from a 0.8B that runs on a laptop to a 27B for a single data-centre GPU.
21
+ - Documents of up to 65,536 tokens, on CUDA and on Apple Silicon through MLX. Each model card says how long a document can get before accuracy drops.
22
+ - Fine-tune on your own labelled examples. A coding-agent skill runs the whole loop on Modal, from finding your questions to serving the result.
23
+ - Deploy your own HTTPS endpoint with one command. It scales to zero when idle.
24
+ - Try it in the browser first: [huggingface.co/spaces/jaredpalmer/kev](https://huggingface.co/spaces/jaredpalmer/kev).
25
+
26
+ ## Models
27
+
28
+ Start with Kev-4B. Move to Kev-9B if you have a bigger GPU, or to Kev-27B if you have an 80 GB GPU and want the most accurate Kev. Use Kev-0.8B when size matters more than accuracy.
29
+
30
+ | Model | Base (license) | Runs on: CUDA | Runs on: Mac (MLX) | Validated context | Held-out datasets: index | Card |
31
+ |---|---|---|---|---|---|---|
32
+ | [Kev-0.8B](https://huggingface.co/jaredpalmer/kev-0.8b) | Qwen3.5-0.8B-Base (Apache-2.0) | L4, any 4 GB GPU | Any Apple Silicon Mac; measured to 65k tokens | 8,192 | 23.3 | [Details](docs/model-cards/kev-0.8b.md) |
33
+ | [Kev-4B](https://huggingface.co/jaredpalmer/kev-4b) | Qwen3.5-4B-Base (Apache-2.0) | L40S, H100 | 32 GB Mac; measured to 65k tokens | 8,192 | 38.0 | [Details](docs/model-cards/kev-4b.md) |
34
+ | [Kev-9B](https://huggingface.co/jaredpalmer/kev-9b) | Qwen3.5-9B-Base (Apache-2.0) | L40S, H100 | 32 GB Mac or larger (expected, not measured) | 8,192 | 41.0 | [Details](docs/model-cards/kev-9b.md) |
35
+ | [Kev-27B](https://huggingface.co/jaredpalmer/kev-27b) | Qwen3.8-27B, post-trained (Apache-2.0) | B200, H200, H100 80 GB | 96–128 GB Mac (expected, not measured) | 65,536 | **52.3** | [Details](docs/model-cards/kev-27b.md) |
36
+ | Jev | Hosted | TypeSafe's API | – | – | 54.0 | – |
37
+
38
+ "Held-out datasets" is the chance-corrected index of the community Decision Index, scored on the test split of `breadth-v1`: 14 public datasets in five areas that no Kev trained on. "Validated context" is the longest document, in tokens, for which accuracy on real contracts (CUAD) stays within 3 points of the same model's accuracy at 8k tokens, at the 95 % lower bound; each model card has the measurement by length.
39
+
40
+ | Model | Accuracy: New Sources | Accuracy: Trained Sources | Brier: New Sources |
41
+ |---|---|---|---|
42
+ | Kev-0.8B | 0.648 / 0.697 | 0.827 / 0.838 | 0.481 / 0.416 |
43
+ | Kev-4B | 0.817 / 0.838 | 0.873 / 0.865 | 0.269 / 0.242 |
44
+ | Kev-9B | 0.820 / 0.852 | 0.874 / 0.873 | 0.289 / 0.217 |
45
+ | Kev-27B | **0.851 / 0.889** | 0.865 / 0.866 | **0.225 / 0.156** |
46
+ | Jev | 0.857 / – | 0.845 / – | 0.211 / – |
47
+
48
+ Each cell is **development / test**. "New sources" means datasets and policy rules Kev never saw during training. It is the closest thing here to your own questions. "Trained sources" means held-out examples from the datasets Kev was trained on. We pick checkpoints using the development sets and read each test set only once per released model. Jev has only been run on the development sets of these two suites. Brier scores the whole probability distribution, not just the top answer; lower is better.
49
+
50
+ On new sources Kev-27B is within a point of Jev (0.851 vs 0.857), and Kev-4B and Kev-9B are within four points. We don't know what Jev was trained on, so this isn't a controlled comparison of the two architectures. [What to Expect](#what-to-expect) says where Kev is as good as Jev and where it isn't.
51
+
52
+ Kev-0.8B, 4B and 9B start from Qwen base models and share one training recipe: a small adapter on a frozen base. Kev-27B starts from Qwen's post-trained release, and we don't know what that was trained on; every one of its weights is fine-tuned, so it ships as 51 GB of full weights rather than an adapter. Each model card has the full recipe, all results, and the earlier versions kept as Hub tags.
53
+
54
+ ## Kev 1.0
55
+
56
+ The four models above are released together as Kev 1.0. Each Hub repo has a `v1.0` tag, so `--run jaredpalmer/kev-4b@v1.0` always loads the same weights, and the GitHub release [`kev-1.0`](https://github.com/jaredpalmer/kev/releases/tag/kev-1.0) has the 0.8B, 4B and 9B checkpoints with SHA-256 checksums. Kev-27B's 51 GB of weights are too large for a release asset and are on the Hub only.
57
+
58
+ | Model | Hub revision of the weights | Temperature | Trained on states up to |
59
+ |---|---|---|---|
60
+ | Kev-0.8B | `9a45d25e` | 2.35 | 7,552 tokens |
61
+ | Kev-4B | `139fdd94` | 2.41 | 7,552 tokens |
62
+ | Kev-9B | `b5d8c18e` (v2) | 2.19 | 7,552 tokens |
63
+ | Kev-27B | `28be62e9` (v2, full weights) | 1.32 | 32,768 tokens |
64
+
65
+ Kev 1.0 trains nothing new. It fixes the checkpoints, cards, evaluation suites and serving code that the next generation of Kev will be compared against. The [release notes](docs/releases/kev-1.0.md) list what changed since the previous family release and what is known not to work well.
66
+
67
+ ## Quick Start
68
+
69
+ ### Try It in the Browser
70
+
71
+ The [Hugging Face Space](https://huggingface.co/spaces/jaredpalmer/kev) runs Kev-4B and Kev-0.8B, with nothing to install.
72
+
73
+ ### Run It Locally
74
+
75
+ You'll need Python 3.12 or 3.13 and [uv](https://docs.astral.sh/uv/). The repo's `.python-version` makes `uv sync` use 3.13; torch has no wheels for 3.14 yet.
76
+
77
+ ```bash
78
+ git clone https://github.com/jaredpalmer/kev.git && cd kev
79
+ uv sync --extra serve
80
+ uv run --extra serve python -m kev.serve --run jaredpalmer/kev-4b --port 8009
81
+ ```
82
+
83
+ This starts Kev-4B on your machine: CUDA or ROCm if you have a GPU, MLX on Apple Silicon. The first run downloads the adapter and the base model. `--run` also accepts a local checkpoint directory or a Hub revision like `jaredpalmer/kev-4b@qwen3`.
84
+
85
+ In another terminal, send it a ticket:
86
+
87
+ ```bash
88
+ curl -s localhost:8009/v1/systemone -H 'content-type: application/json' -d '{
89
+ "state": "Shoes arrived two weeks late and in the wrong size. Also I see two charges on my card.",
90
+ "model": "kev-latest",
91
+ "questions": {
92
+ "department": {"type": "choice", "instructions": "Which team should handle this?",
93
+ "criteria": {"returns": "Exchanges, refunds, wrong or damaged items",
94
+ "shipping": "Delivery status, delays, lost packages",
95
+ "billing": "Charges, invoices, payment problems"}},
96
+ "escalate": {"type": "noul", "instructions": "Does this need urgent human attention?"},
97
+ "frustration": {"type": "score", "instructions": "How frustrated is the customer?",
98
+ "criteria": ["Calm", "Frustrated", "Very angry"]}
99
+ }}'
100
+ ```
101
+
102
+ Example response from Kev-4B, running in bf16 on an Apple M5:
103
+
104
+ ```json
105
+ {
106
+ "model": "kev-latest",
107
+ "answers": {
108
+ "department": { "type": "choice", "choice": "returns", "confidence": 0.21,
109
+ "probabilities": { "returns": 0.47, "shipping": 0.28, "billing": 0.25 } },
110
+ "escalate": { "type": "noul", "noul": 0.93 },
111
+ "frustration": { "type": "score", "score": 1.44, "confidence": 0.34,
112
+ "legend": { "0": "Calm", "1": "Frustrated", "2": "Very angry" },
113
+ "probabilities": { "0": 0.00, "1": 0.56, "2": 0.44 } }
114
+ },
115
+ "usage": { "input_tokens": 101, "output_tokens": 161 },
116
+ "latency_ms": 495
117
+ }
118
+ ```
119
+
120
+ The ticket mentions a return, a late delivery and a billing problem, and the department probabilities say so. That's why Kev returns probabilities instead of a single label: your code can route the confident cases and send the rest to a person.
121
+
122
+ ### Use It From Python
123
+
124
+ If you already call Jev, point your client at Kev and keep the rest of your code. The TypeSafe SDK is included in `uv sync --extra serve`:
125
+
126
+ ```python
127
+ from typesafe_sdk import Choice, Noul, Score, TypeSafeClient
128
+
129
+ client = TypeSafeClient(
130
+ api_key="local",
131
+ base_url="http://127.0.0.1:8009",
132
+ model="kev-latest",
133
+ )
134
+ response = client.system_one(
135
+ state="I was charged twice. Please fix this ASAP.",
136
+ questions={
137
+ "billing": Noul(instructions="Is this ticket about billing?"),
138
+ "tone": Choice(
139
+ instructions="What is the customer's tone?",
140
+ criteria={"calm": None, "frustrated": None, "angry": None},
141
+ ),
142
+ "urgency": Score(
143
+ instructions="How urgent is this ticket?",
144
+ criteria=["can wait", "this week", "today"],
145
+ ),
146
+ },
147
+ )
148
+ print(response.nouls["billing"].noul)
149
+ print(response.choices["tone"].choice)
150
+ print(response.scores["urgency"].score)
151
+ ```
152
+
153
+ ## Fine-Tune on Your Own Data
154
+
155
+ The released models were trained on public datasets and generated policy examples. If your questions look different, like your own routing categories, your own escalation rules or another language, a short fine-tune usually helps more than any prompt change. It also fits the temperature to your data, so the confidence you set thresholds on is measured on your own labels.
156
+
157
+ What to expect: on an example support workload (three questions, 1,050 generated records, 15 minutes on an H100), fine-tuning took Kev-4B from 67.7% to 73.6% accuracy, and from automating 34% of decisions at a 5% error budget to 48% ([details](skills/kev-finetune/README.md#why-fine-tune-at-all)). On real data, one epoch on 5,219 labelled consumer-finance complaints took Kev-4B from 0.804 to 0.904 accuracy on complaints it had never seen. Gains like these are in distribution: they tell you how well Kev learns your task, not how it does on everything else. Size your dataset first. With 400 records, the gain on the example workload was inside the noise.
158
+
159
+ ### With a Coding Agent
160
+
161
+ ```bash
162
+ npx skills add jaredpalmer/kev@kev-finetune
163
+ ```
164
+
165
+ Then ask your agent to "fine-tune Kev on my support tickets". The [`kev-finetune` skill](skills/kev-finetune/) interviews you, finds the questions your code already asks Jev or TypeSafe, converts the labels you have or generates enough with any LLM to measure a gain, fine-tunes from a released checkpoint on Modal, fits the temperature on a held-out slice, scores the result against the untouched model, deploys an endpoint, and tears everything down at the end. You don't need a local GPU or a clone of this repo. A Kev-4B training run costs about $1 on an H100.
166
+
167
+ ### By Hand
168
+
169
+ The skill's [README](skills/kev-finetune/README.md) is the same recipe for people: six short standard-library scripts and one Modal app. To train from this repo instead, put your examples in a JSONL file, one request per line. It's the same shape as an API request, plus a `label` on every question:
170
+
171
+ ```jsonl
172
+ {"state": {"subject": "Charged twice", "body": "I see two charges for order #4411. Please refund one."},
173
+ "questions": {
174
+ "team": {"type": "choice", "instructions": "Which team should handle this ticket?",
175
+ "criteria": {"billing": "Payments and refunds", "shipping": "Delivery problems", "access": "Login and account access"}, "label": "billing"},
176
+ "angry": {"type": "noul", "instructions": "Is the customer angry?", "label": false},
177
+ "priority": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["low", "normal", "high"], "label": 1}}}
178
+ ```
179
+
180
+ For `choice` the label is the option name, for `noul` it's `true` or `false`, and for `score` it's the level's position starting at 0. Keep 10–20% of the file aside for evaluation.
181
+
182
+ Then start from a released checkpoint with `--init_from`:
183
+
184
+ ```bash
185
+ uv run python -m kev.train --data train.jsonl --base Qwen/Qwen3.5-4B-Base --init_from jaredpalmer/kev-4b \
186
+ --epochs 2 --lr 2e-5 --batch 1 --accum 8 --dtype bf16 --checkpointing 1 --device cuda --out runs/mine
187
+
188
+ uv run python -m kev.benchmark --run runs/mine --data heldout.jsonl --out runs/mine-eval
189
+ uv run --extra serve python -m kev.serve --run runs/mine --port 8009
190
+ ```
191
+
192
+ `--init_from` loads the adapter and pointer head from the released model before training, so you keep what Kev already knows and add your domain on top. Starting from the base model instead throws that away: in one user's test on 836 support-tool decisions, a fine-tune from the base scored 0.33 on Kev's own evaluation set, against 0.84 for the released model; the same data with `--init_from` kept 0.83 there and reached 0.88 on the new domain. Use a smaller learning rate than the from-scratch recipe (`2e-5` is a good start), and pick `--base` to match the checkpoint you start from; the trainer checks that the base, revision, LoRA rank and head size agree before it loads anything.
193
+
194
+ `--batch 1 --accum 8` in bf16 fits the 0.8B model on a 4 GB GPU. The benchmark reports accuracy, Brier score and calibration per question type, so you can see which of your questions the fine-tune helped. The checkpoint you started from is recorded in `runs/mine/training_config.json`. On a Mac, run one training job at a time; two jobs on the same Apple GPU are much slower.
195
+
196
+ ## Deploy Your Own Endpoint
197
+
198
+ To get an HTTPS endpoint instead of a local server, you don't need this repo, just a [Modal](https://modal.com) account:
199
+
200
+ ```bash
201
+ pip install modal && modal setup
202
+ curl -LO https://raw.githubusercontent.com/jaredpalmer/kev/main/skills/kev-deploy/scripts/kev_serve.py
203
+ KEV_API_KEY=$(openssl rand -hex 24) modal deploy kev_serve.py
204
+ ```
205
+
206
+ That serves Kev-4B on an L40S at `https://<your-workspace>--kev-api.modal.run`, with the same API as above behind `Authorization: Bearer <key>`. It scales to zero when idle, so an unused endpoint costs nothing. The first request after idle waits about 35 seconds for a container to start. `KEV_MODEL=jaredpalmer/kev-9b` serves another model on the GPU that suits it; Kev-27B goes to a B200, falling back to an H200 or H100. If you use a coding agent, `npx skills add jaredpalmer/kev@kev-deploy` does the same and wires the URL into your code. [skills/kev-deploy](skills/kev-deploy/) has the GPU and cost table.
207
+
208
+ A model you fine-tuned with the `kev-finetune` skill deploys the same way from its own Modal app (`KEV_SERVE_SECRET=kev-serve-key KEV_SERVE_RUN=<run> modal deploy scripts/kev_modal.py`; see [its deploy guide](skills/kev-finetune/references/deploy.md)). To host Kev on your own machines instead, run `kev.serve` from [Run It Locally](#run-it-locally) on a GPU box with `--host 0.0.0.0` and put it behind your own proxy; [Serving Performance](#serving-performance) says which GPU to pick.
209
+
210
+ ## What to Expect
211
+
212
+ **Accuracy.** Kev-27B is within three points of Jev, or ahead of it, on 9 of the 11 new-source categories in the chart below. Kev-4B and Kev-9B are about as close on classification-shaped sources like routing, entailment and science questions. Knowledge questions depend mostly on the base model: on MMLU Kev-9B scores 0.73 and Kev-27B matches Jev at 0.90, but on the harder MMLU-Pro Kev-27B scores 0.675 against Jev's 0.840. The smaller models also trail on day-precision date arithmetic.
213
+
214
+ ![Accuracy by source for Kev and Jev](docs/kev-family.png)
215
+
216
+ **Confidence.** Each checkpoint ships with a fitted temperature, so its probabilities are calibrated by default. As served, Kev-9B puts at least 0.9 probability on a wrong answer for 2.4% of new-source questions, against Jev's 3.7%. Jev still ranks its answers better: at a 5% error budget, Kev-4B, 9B and 27B can automate 0.52–0.69 of new-source decisions, Kev-0.8B 0.14 and Jev 0.70. Check a threshold on your own data before you rely on it.
217
+
218
+ **Speed.** Kev-4B answers six questions about a new short text in 18.1 ms of model time on an H100 and 41.5 ms on an L40S, and a container serves around 101 requests per second on an H100. On an Apple M5, Kev-4B takes 721 ms for five questions, or 136 ms when the text repeats and comes from the cache. [Serving Performance](#serving-performance) has every GPU and batch size.
219
+
220
+ **Length.** Kev-0.8B, 4B and 9B trained mostly on states of up to 384 tokens, with longer ones in their document and skill fine-tunes (up to 7,552 tokens), Kev-27B on states of up to 32,768. The server accepts states of up to 65,536 tokens, and 8,192 more for each question, and refuses a longer one with a 422 instead of cutting it. How far past its training length each model stays accurate is the "Validated context" column in [Models](#models). For Kev-0.8B, 4B and 9B that is 8,192 tokens: at 16k the measurement on real contracts can no longer rule out a drop of more than 3 points, and at 32k all three are measurably less accurate than at 8k. Kev-27B holds to the 65,536-token limit. On real contracts of up to 64k tokens (CUAD) Kev-27B scores 0.874, and its confidence there is less reliable than on short text; its [model card](docs/model-cards/kev-27b.md) has the numbers by length.
221
+
222
+ ## Playground
223
+
224
+ With the server running, open another terminal. You'll need Node 20.9+:
225
+
226
+ ```bash
227
+ cd playground
228
+ npm install
229
+ npm run dev -- -p 3001
230
+ ```
231
+
232
+ Open [localhost:3001](http://localhost:3001), load a preset, and edit the text and questions. Press `⌘↵` to run it. "Packed vs separate" compares asking all questions at once with asking them one at a time. "Permute" runs a Choice question with six option orders. There are also presets for testing question isolation and fake delimiter tokens.
233
+
234
+ ![Kev playground](docs/playground.png)
235
+
236
+ There's a [chess demo](http://localhost:3001/chess), too. The board is the input, legal moves are Choice options, and a Score question rates the position. You can play against Kev or let it play itself. Games are saved in `localStorage`.
237
+
238
+ ## API
239
+
240
+ ### `POST /v1/systemone`
241
+
242
+ `state` is the text to evaluate. Each question has instructions and, where needed, a set of answers to choose from.
243
+
244
+ ```jsonc
245
+ {
246
+ "state": "…", // string | object | array — the content to evaluate
247
+ "model": "kev-latest",
248
+ "questions": {
249
+ "<id>": { // you choose the id; the model never sees it
250
+ "type": "noul" | "choice" | "score",
251
+ "instructions": "…", // string | object | array, optional
252
+ "criteria": … // noul: {true?, false?} choice: {option: description|null} score: [level, …]
253
+ }
254
+ }
255
+ }
256
+ ```
257
+
258
+ | Type | Criteria | Answer |
259
+ |---|---|---|
260
+ | `noul` | Optional descriptions for `true` and `false` | `noul`: probability of yes |
261
+ | `choice` | 1–255 option names, each with a description or `null` | `choice`: most likely option; `probabilities` and `confidence` |
262
+ | `score` | 1–255 descriptions, ordered from lowest to highest | `score`: mean level index, starting at 0; `legend`, `probabilities`, and `confidence` |
263
+
264
+ For Choice with `K > 1` options, confidence is `(p_max − 1/K) / (1 − 1/K)`. A single option has confidence 1. Score confidence is `max(0, 1 − E|level − mode| / D)`: `mode` is the most likely level and `D` is the mean distance of a uniform distribution over the levels from its middle (2/3 for three levels), so all probability on one level gives 1 and a uniform or wider spread gives 0. Both formulas are the ones in TypeSafe's reference adapter ([`system-one-adapter`](https://github.com/typesafe-ai/system-one-adapter-python) 0.2.1). Neither field is a measured accuracy rate.
265
+
266
+ Objects and arrays are converted to labeled text. Delimiter-like strings in user input are escaped before tokenization. Invalid requests return `422`, and so does a state longer than 65,536 tokens: the server never drops part of a document silently, and the error gives the state's token count and the limit. `usage.output_tokens` counts tokens in the serialized answers, not generated tokens.
267
+
268
+ | Method | Path | Purpose |
269
+ |---|---|---|
270
+ | `GET` | `/v1/models` | Model cards (`name`, `description`, `release_date`) plus the loaded checkpoint's details |
271
+ | `POST` | `/v1/systemone/permute` | Run one Choice question with different option orders (`n_perm` 1 to 64, default 6) |
272
+ | `POST` | `/v1/systemone/separate` | Run each question in its own forward pass |
273
+
274
+ A request may carry any number of questions. The server runs them a token budget at a time (one maximal row of 16,384 tokens per forward pass, counting the cached document once per question in that pass), so memory does not grow with the question count and the answers do not depend on the split. Every response carries an `x-typesafe-request-id` header. The server binds to `127.0.0.1` (`--host 0.0.0.0` to accept other machines) and is open by default; set `KEV_API_KEY` to require `Authorization: Bearer <key>` on `/v1/*`, as the TypeSafe clients always send it.
275
+
276
+ | Variable | Effect |
277
+ |---|---|
278
+ | `KEV_TEMPERATURE=1.0` | Return raw probabilities instead of the calibrated ones |
279
+ | `KEV_DATE_FACTS=1` | Append the number of days between any two dates in the state (see [Benchmarks](#benchmarks)) |
280
+ | `KEV_TRUNCATE_STATES=1` | Read the first 65,536 tokens of a longer state instead of refusing it; every response then has `truncated` and `usage.state_tokens` / `state_tokens_used` |
281
+ | `KEV_DTYPE=fp32` | Serve the exact fp32 path the evaluations use (bf16 is the default on GPUs) |
282
+ | `KEV_API_KEY` | Require a bearer key |
283
+
284
+ ## How It Works
285
+
286
+ Each checkpoint is a rank-16 LoRA adapter and a small pointer head on a Qwen base model. On an attention-only base (Qwen3), the state and questions go into one token sequence:
287
+
288
+ ```text
289
+ <state> …state…
290
+ <q> instructions <opt> option 1 </opt> <opt> option 2 </opt> … <decide>
291
+ <q> instructions <opt> option 1 </opt> <opt> option 2 </opt> … <decide>
292
+ ```
293
+
294
+ The attention mask lets a token read the state and its own question, but not other questions or future tokens. Each question's position IDs restart just after the state. This lets the model process the state once and answer each question independently.
295
+
296
+ Qwen3.5 and Qwen3.8 mix attention layers with Gated DeltaNet layers, which are recurrent and ignore attention masks. For those models, which is every current Kev, each question runs as its own row: the state followed by that question, with the same positions as above. The rows are independent, so isolation is exact, and the server and `DecisionModel.probs()` compute the state once and reuse its cache for every row. `forward()`, which `kev.benchmark` scores and every published number comes from, keeps the plain rows and runs the state once per question; the two agree to fp32 rounding. On attention-only models the rows and the mask above give identical probabilities (`tests/test_model.py`).
297
+
298
+ Kev-27B uses the same design on `Qwen/Qwen3.8-27B`, with two differences. Its base is Qwen's post-trained release rather than a `-Base` checkpoint, and we don't know what it was post-trained on. And every backbone weight is trained, not just an adapter, and kept in bf16, so the checkpoint is the whole model: 51 GB of bf16 weights plus the pointer head. It serves in bf16 only (about 66 GB resident with the serving buffers), which is why it needs an 80 GB card. On Apple Silicon the MLX backend loads those weights as they are, with no merge (see [Serving Performance](#serving-performance)); we expect that to fit a 96–128 GB Mac but haven't measured it. Its served probabilities stay within 0.022 of the evaluation path on an H200 (`runs/serving-27b-r23`).
299
+
300
+ The pointer head scores each option's `</opt>` hidden state against the question's `<decide>` hidden state. A softmax turns those scores into probabilities. Because `<decide>` comes last, it can attend to the full option list.
301
+
302
+ Training uses cross-entropy on the correct answer. The adapter and head are trained together; the rest of the base weights stay fixed (Kev-27B trains them all). Training examples and API requests use the same text format. No Jev outputs were used for training.
303
+
304
+ Asking questions together or separately produces probabilities within 4e-6 in the fp32 tests. This does **not** mean option order is irrelevant: options within a question can still affect one another. See [the model code](kev/model.py) and [parity tests](tests/test_model.py).
305
+
306
+ ## Training
307
+
308
+ The released models share one base training set, `decision-v7`: 10,000 examples from ten public datasets, 896 generated policy examples, and 1,680 examples from 60 generated rule structures. Kev-0.8B, 4B and 9B train on it for two epochs with LoRA rank 16 and cross-entropy. The learning rate is `1e-4` for 0.8B and `5e-5` for 4B and 9B. On these hybrid bases the adapter covers the attention, MLP and DeltaNet projections; `kev.train` picks the right targets from the model config.
309
+
310
+ Kev-0.8B, 4B and 9B then get short follow-up fine-tunes from their released checkpoints, through the same `--init_from` path you'd use for your own data: generated cases that state day counts or have the deciding evidence removed (all three), then real documents and generated skill data (all three; Kev-9B since v2, 2026-09-30). Kev-27B is trained differently. Every weight of the base is fine-tuned for one epoch on eight H200s (`--full_ft 1`, learning rate `2e-6`) on a 145,840-record corpus: Kev's own data, the document, skill and developer-tooling suites, public datasets, licensed task families and generated long-document, tool-routing, agent-log and guardrail records, with states of up to 32,768 tokens. The result is then averaged with the earlier adapter-trained Kev-27B, 0.85 to 0.15. The model cards list every stage with its data and cost.
311
+
312
+ ```bash
313
+ # sanity run, ~1 minute
314
+ uv run python -m kev.train --n_per_source 40 --accum 4 --out runs/smoke
315
+
316
+ # the first stage of Kev-0.8B (~20 min on one H100; the Mac path works but is slow for Qwen3.5 bases)
317
+ uv run python -m kev.train --suite evals/v7/decision-v7 --base Qwen/Qwen3.5-0.8B-Base --base_revision dc7cdfe2ee4154fa7e30f5b51ca41bfa40174e68 \
318
+ --epochs 2 --lr 1e-4 --batch 8 --dtype bf16 --p_none_pair 0.25 --device cuda --out runs/kev-0.8b
319
+
320
+ # the first stage of Kev-4B (one H100 via Modal, ~1 h; see below). Swap in Qwen/Qwen3-4B-Base for the previous generation.
321
+ uv run python -m kev.train --suite evals/v7/decision-v7 --base Qwen/Qwen3.5-4B-Base --base_revision 1001bb4d826a52d1f399e183466143f4da7b741b \
322
+ --epochs 2 --lr 5e-5 --batch 4 --accum 2 --dtype bf16 --checkpointing 1 --p_none_pair 0.25 --device cuda --out runs/kev-4b
323
+ ```
324
+
325
+ Use `uv run python -m kev.train --help` for all training options. The released models don't use the optional `--perm_kl` or `--ord_w` losses. [PLAN.md](PLAN.md) records what was tried, what helped, and what didn't.
326
+
327
+ ### Modal
328
+
329
+ Each trial gets its own H100. The study keeps running if you disconnect, and you can download the results when it finishes:
330
+
331
+ ```bash
332
+ uv run modal token new # once; opens the browser
333
+ KEV_GPU=T4 uv run modal run modal_app.py::smoke # end-to-end check, ~1 minute of GPU
334
+
335
+ uv run modal deploy modal_app.py # once; studies run on the deployed app and survive disconnects
336
+ uv run modal run modal_app.py::study \
337
+ --suite evals/v7/decision-v7 --plan experiments/v7-final.json \
338
+ --name my-study --transfer evals/v4/transfer-v4 --budget 30 --timeout 7200
339
+ uv run modal run modal_app.py::pull --name my-study # results -> runs/my-study, ranked
340
+ ```
341
+
342
+ [Study plans](experiments/v7-final.json) list training settings. Each trial saves the settings, code hashes, dataset hashes, and results. Choose models using the development results, not the locked test. After choosing a final candidate, you can read its test results once:
343
+
344
+ ```bash
345
+ uv run modal run modal_app.py::locked_test --trial my-study/00-trial-0 --name my-candidate # one read, ever
346
+ ```
347
+
348
+ ## Benchmarks
349
+
350
+ The evaluation data under `evals/` is frozen: dataset versions and file checksums are recorded in each manifest. Large files are downloaded from [the Hub mirror](https://huggingface.co/datasets/jaredpalmer/kev-suites) and checked against those hashes. Every model in the tables above is scored on the same items. The numbers in this README and the model cards are checked in CI against the committed reports they come from (`docs/claims.json`, `uv run python scripts/verify_claims.py`).
351
+
352
+ | Suite | What it measures |
353
+ |---|---|
354
+ | `decision-v7` | Held-out examples from the ten training datasets, the generated policies and the rule structures ("trained sources") |
355
+ | `transfer-v4` | 764 records from datasets and policy and rule types Kev never trained on: QNLI, SciQ, PAWS, MMLU, Emotion, TweetEval, held-out policies and rules ("new sources") |
356
+ | `transfer-v9` | `transfer-v4` plus 10-way MMLU-Pro, records buried in unrelated text, and "unknowable" records whose deciding evidence was removed |
357
+
358
+ ```bash
359
+ uv run python -m kev.benchmark --run jaredpalmer/kev-4b --suite evals/v4/transfer-v4 --out runs/my-eval # new sources
360
+ uv run python -m kev.benchmark --run jaredpalmer/kev-4b --suite evals/v9/transfer-v9 --out runs/my-eval-v9 # + MMLU-Pro, buried states, unknowable items
361
+ uv run python -m kev.benchmark --run jaredpalmer/kev-4b --suite evals/v7/decision-v7 --out runs/my-eval-id # trained sources
362
+ uv run python -m kev.benchmark --remote http://127.0.0.1:8009 --suite evals/v4/transfer-v4 --out runs/my-remote # any System One endpoint, Jev included
363
+ ```
364
+
365
+ These commands use development data. Test data requires `--allow-test`. The benchmark reports accuracy, Brier score, calibration error, the share of decisions you could automate at a 5% error budget, option-order changes, and question isolation. On the unknowable records it reports how often the model still answers with at least 0.9 confidence (Kev-9B 0%, Jev 9%). Published accuracy numbers use fp32 evaluation, not the bf16 serving path. `kev.jev` runs the same questions against Jev through Vercel AI Gateway, and `kev.compare` compares two saved runs with paired bootstrap confidence intervals.
366
+
367
+ **Calibration.** Each checkpoint stores a temperature, and the pointer head applies it when the model is loaded. Kev-4B (2.41) and Kev-0.8B (2.35) fitted theirs on their in-distribution development sets; Kev-27B (1.32) and Kev-9B (2.19) fitted theirs on held-out datasets they never trained on. A refit of the two smaller models on those held-out datasets was tested and kept neither: it did not improve Kev-4B and made Kev-0.8B worse calibrated on its document and skill suites (the model cards have the numbers). A temperature never changes which answer wins. On new sources it takes Kev-9B's calibration error from 0.103 to 0.041 and its confident errors (wrong answers with probability ≥ 0.9) from 8.2% to 2.4%, below Jev's 3.7%. The accuracy numbers above are the same either way; the Brier numbers are for the raw probabilities. `scripts/calibrate_checkpoint.py` also reports an out-of-fold estimate, so the in-sample fit can be checked against records it didn't see.
368
+
369
+ **Dates.** Kev can't subtract dates reliably, but it can use a day count it's given. `KEV_DATE_FACTS=1` appends one sentence per pair of dates in the state ("June 26, 2026 is 8 days before July 4, 2026"). On the deadline policy questions this takes Kev-9B from 0.80 to 0.90 (Jev 0.93). None of the tables use it.
370
+
371
+ **Other people's test sets.** `evals/external/` holds test sets from other projects, converted to this format, with their published live Jev results. Some were scored on earlier versions of the Kev weights, which the Kev column names. Three were removed because they cannot serve as a gate, and the model cards keep the numbers their releases were decided on: scienthoon's synthetic support tickets on 2026-09-27 (templated text; one of its three questions depends on a rule the text does not state), and on 2026-09-30 WANLI (`wanli-v1`, `wanli-v2`: a quarter of the pairs are ones WANLI's two annotators labelled differently, with the gold set to one of them) and TypeSafe's public evals (`typesafe-v1`: the gold is the averaged answer of two closed frontier models, and on 89 questions it cannot tell checkpoints apart).
372
+
373
+ | Suite | What it is | Jev | Kev |
374
+ |---|---|---|---|
375
+ | [SemIf](https://github.com/TheoLeeCJ/SemIf) | 144 authored decisions | 0.965 | 0.917 (Kev-9B at `v7-base`) |
376
+
377
+ SemIf's labels hold up, but it is close to saturated: every Kev-27B checkpoint answers 130 of the 144 correctly, so it is a sanity check, not a way to rank models.
378
+
379
+ ## Serving Performance
380
+
381
+ Pick the GPU by the model:
382
+
383
+ | Model | GPU ($/h) | 6 questions, short text | 5 questions, 2,200-token text | Requests/s, 64 clients |
384
+ |---|---|---|---|---|
385
+ | Kev-0.8B | L4 (0.80) | 22.7 / 16.1 ms | 108.6 / 32.3 ms | 62.8 |
386
+ | Kev-4B | L40S (1.95) | 41.5 / 27.7 ms | 145.2 / 43.0 ms | 51.4 |
387
+ | Kev-4B | H100 (3.95) | 18.1 / 12.9 ms | 89.4 / 22.5 ms | 100.8 |
388
+ | Kev-9B | L40S (1.95) | 66.4 / 42.7 ms | 235.6 / 57.5 ms | 32.7 |
389
+ | Kev-9B | H100 (3.95) | 24.0 / 16.6 ms | 88.5 / 26.4 ms | 79.5 |
390
+ | Kev-27B | B200 (6.25) | 46.5 / 32.2 ms | 178.0 / 52.1 ms | 44.2 |
391
+ | Kev-27B | H200 (4.54) | 67.2 / 50.0 ms | 274.8 / 73.8 ms | 28.6 |
392
+ | Kev-27B | H100 (3.95) | 75.0 / 52.0 ms | 277.5 / 79.3 ms | 28.9 |
393
+
394
+ Times are model time per request (the `latency_ms` the API returns), median of 20, for a new text / the same text again. The server caches the text, so asking more questions about a document you've already sent only pays for the questions. Requests per second are for 64 concurrent clients sending six questions about a new short text each; the server batches them. Kev-27B's B200 and H100 rows were measured on its previous version, the same architecture served in bf16 (`runs/fused-27b-*`); the H200 row is the current checkpoint (`runs/serving-27b-r23`). Network time is extra: about 65 ms per round trip through a Modal web endpoint in the same region.
395
+
396
+ An L4 is enough for Kev-0.8B but too slow for Kev-4B. The A100 is slower than the L40S here and costs more. Kev-9B needs about 17 GB of GPU memory and Kev-27B 51 GB of weights (about 66 GB with the batching buffers); under load Kev-27B is compute-bound, and a B200, H200 or H100 costs about the same per request. On CUDA, install `flash-linear-attention` for the Qwen3.5 models (`kev_serve.py` and the Modal images already do).
397
+
398
+ On Apple Silicon, `uv sync --extra serve` installs [MLX](https://github.com/ml-explore/mlx-lm) and the server uses it automatically. Five questions about a ~270-token text on an M5 (32 GB):
399
+
400
+ | Model | New text | Same text again |
401
+ |---|---|---|
402
+ | Kev-0.8B | 149 ms | 28 ms |
403
+ | Kev-4B | 721 ms | 136 ms |
404
+
405
+ Long documents are read into the cache 1,024 tokens at a time, so memory stays close to the weights. With a 65,000-token document, Kev-0.8B takes 21.2 s the first time and 202 ms after that, at a 3.8 GB peak, and Kev-4B 84.5 s and 716 ms at 13.0 GB (`runs/mlx-long-states`; the model cards have every length). Kev-9B has not been measured this way yet.
406
+
407
+ Adapter checkpoints are folded into the base as they load, which briefly holds a second copy of the weights. Full-weight checkpoints like Kev-27B load as saved, with nothing merged, so loading needs only the weights. We checked this on Kev-4B written out as full bf16 weights: loading peaked at 8.4 GB for 8.4 GB of weights, against 15.9 GB for the adapter path. Its answers matched the adapter path exactly once both hold the same bf16 values, and stayed within 0.015 of the fp32 path on 60 questions (`runs/mlx-full-4b`). Kev-27B's weights are 51 GB. By the same measurements it needs about 51 GB plus working memory, so a 64 GB Mac is borderline and a 96–128 GB Mac should fit. We haven't run it on a Mac that large yet. Kev-27B's first version, an adapter, did run this way on a 128 GB M5 Max, matching the published accuracy (thanks to Sean Connelly, [#175](https://github.com/jaredpalmer/kev/pull/175)).
408
+
409
+ The server runs in bf16 on GPUs and Macs. Its probabilities differ from the fp32 path the published evaluations use by at most about 0.03 on a GPU and 0.05 on a Mac, and the top answer changes on about one question in 300. Set `KEV_DTYPE=fp32` for the exact path. `/v1/models` reports the backend and precision in use. `uv run modal run modal_app.py::serving --run jaredpalmer/kev-4b --gpu L40S --name <name>` measures a row of the table on your own account (the rows above: `runs/serve-*`, `runs/grouping-4b-h100`, `runs/fused-27b-*`, `runs/serving-27b-r23`).
410
+
411
+ ## Limitations
412
+
413
+ - Calibration is a single temperature. It can't reorder confidences, so the share of new-source decisions you can automate at a 5% error budget (0.52–0.69 for Kev-4B, 9B and 27B) is still below Jev's 0.70. Test a probability threshold on your own data before you rely on it.
414
+ - Knowledge questions are set by the base model. MMLU is 0.73 for Kev-9B against Jev's 0.90, and MMLU-Pro 0.59 against 0.84.
415
+ - Fine-tuning can make the base model worse at individual tasks. Date arithmetic was the clearest case ([issue #8](https://github.com/jaredpalmer/kev/issues/8)); training on stated day counts plus `KEV_DATE_FACTS=1` recovers it.
416
+ - Changing option order can change an answer. Question isolation doesn't prevent this.
417
+ - Kev-0.8B, 4B and 9B trained mostly on at most 384 state tokens and 1,024 tokens for the state plus one question (their document and skill fine-tunes on states of up to 7,552 tokens), Kev-27B on states of up to 32,768 tokens. Serving allows a 65,536-token state; the validated context length of each model is in [Models](#models).
418
+ - On a Mac, answers take hundreds of milliseconds, not tens. Kev-27B needs an 80 GB GPU. On a Mac it needs about 51 GB plus working memory; we expect a 96–128 GB Mac to fit it but haven't measured one.
419
+ - Kev-27B starts from a post-trained model whose training data we don't know.
420
+
421
+ ## Development
422
+
423
+ ```bash
424
+ uv run --extra serve python -m pytest tests/test_unit.py tests/test_research.py tests/test_generators.py tests/test_conventions.py \
425
+ tests/test_documents_tools.py tests/test_hard_v1.py tests/test_devtools_v1.py tests/test_breadth_v1.py tests/test_rounds.py tests/test_skill_scripts.py -q # no weights, no server; what CI runs
426
+ KEV_BASE_URL=http://127.0.0.1:8009 uv run --extra serve python -m pytest tests/test_api.py -q # against a running server
427
+ cd playground && npm run lint && npx next typegen && npx tsc --noEmit -p .
428
+ ```
429
+
430
+ The API tests run TypeSafe's example requests and the official SDK against your local server. [PLAN.md](PLAN.md) is the research plan: what we have learned, the rules every experiment follows, and one line per round. The full log (every experiment, the criteria set before it ran, and how it came out) is at the git tag `research-archive-2026-09-24`.
431
+
432
+ <details>
433
+ <summary>Previous generation (Qwen3) and the prototype</summary>
434
+
435
+ The first Kev family used Qwen3 bases with the same data and settings. Those weights stay published and run on plain PyTorch on a Mac, but they are no longer developed.
436
+
437
+ | Model | Base | Accuracy: Trained Sources | Accuracy: New Sources | Brier: New Sources | Model Card |
438
+ |---|---|---|---|---|---|
439
+ | Kev-0.6B (Qwen3) — `jaredpalmer/kev-0.6b` | Qwen3-0.6B-Base | 0.801 / 0.808 | 0.620 / 0.642 | 0.536 / 0.483 | [Details](docs/model-cards/kev-0.6b-qwen3.md) |
440
+ | Kev-4B (Qwen3) — `jaredpalmer/kev-4b@qwen3` | Qwen3-4B-Base | 0.854 / 0.856 | 0.790 / 0.806 | 0.328 / 0.294 | [Details](docs/model-cards/kev-4b-qwen3.md) |
441
+ | Kev-8B (Qwen3) — `jaredpalmer/kev-8b` | Qwen3-8B-Base | 0.863 / 0.870 | 0.796 / 0.780 | 0.337 / 0.327 | [Details](docs/model-cards/kev-8b-qwen3.md) |
442
+
443
+ The original [Kev-0.5B](https://huggingface.co/jaredpalmer/kev-0.5b) used Qwen2.5-0.5B and is kept for reference; see its [model card](docs/model-cards/kev-0.5b.md).
444
+
445
+ </details>
446
+
447
+ <details>
448
+ <summary>Troubleshooting</summary>
449
+
450
+ - If MPS runs out of memory during training, check that you're running only one job. Don't enable `output_hidden_states` or add tokens with peft's `trainable_token_indices`; both have caused memory problems here.
451
+ - If the playground loads but buttons don't work, use `localhost:3001`. Next.js checks development hostnames. Other hosts need an entry in `allowedDevOrigins` in `playground/next.config.ts`.
452
+ - If dataset loading reports `Dataset scripts are no longer supported`, use `legacy-datasets/banking77`. This repo already uses it.
453
+
454
+ </details>
455
+
456
+ ## Authors
457
+
458
+ - Jared Palmer ([@jaredpalmer](https://github.com/jaredpalmer))
459
+
460
+ Built with [Devin](https://devin.ai). Thanks to [Archer Hume](https://archerhume.com/posts/jevs-architecture-unmasked) for the architecture write-up, [TypeSafe](https://docs.typesafe.ai/api) for the API design, and [Qwen](https://huggingface.co/Qwen/Qwen3.5-9B-Base) for the base models.
461
+
462
+ Related work: [Hydragen](https://arxiv.org/abs/2402.05099), [DeFT](https://arxiv.org/abs/2404.00242), [FIRST](https://arxiv.org/abs/2406.15657).
463
+
464
+ ## License
465
+
466
+ [Apache-2.0](LICENSE). The Qwen3, Qwen3.5 and Qwen3.8 base models are also Apache-2.0. Training datasets have their own licenses; see the [model cards](docs/model-cards/).
runtime/docs/packed-mlx.md ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Standalone Packed MLX Models
2
+
3
+ The packed checkpoint runs on Apple Silicon using MLX. It carries the entire
4
+ compressed backbone, compensation adapter, pointer head, temperature, model
5
+ configuration and tokenizer. Loading and inference never read the original
6
+ backbone or the original PEFT adapter, and work with an empty Hugging Face cache
7
+ and networking disabled.
8
+
9
+ Ternary weights stay in MLX's native 2-bit affine format; selected linear layers
10
+ and the embedding stay 4-bit. CLoQ is applied as two small matrix products,
11
+ without building its dense delta. The embedding dequantizes only requested
12
+ token rows. Activations are bf16. The constructor discards lazy random dense
13
+ parameters before evaluating any model parameters.
14
+
15
+ ## Run
16
+
17
+ From the Kev repository root, with the existing MLX environment:
18
+
19
+ ```sh
20
+ HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-cache \
21
+ KEV_BACKEND=mlx KEV_PREFIX_CACHE=1 KEV_PREFIX_MAX_TOKENS=4096 \
22
+ .venv/bin/python -m kev.serve \
23
+ --run runs/ternary-rotation-20261005/standalone/kev-4b-T-cloq16-mlx \
24
+ --host 127.0.0.1 --port 8010
25
+ ```
26
+
27
+ The server accepts the existing `POST /v1/systemone` API:
28
+
29
+ ```sh
30
+ curl http://127.0.0.1:8010/v1/systemone \
31
+ -H 'Content-Type: application/json' \
32
+ -d '{"state":"I was charged twice for order 8812.","questions":{"team":{"type":"choice","instructions":"Which team should handle this ticket?","criteria":{"billing":"Charges and payments","shipping":"Delivery and delays"}}}}'
33
+ ```
34
+
35
+ Library callers use the usual checkpoint interface:
36
+
37
+ ```python
38
+ from kev.checkpoint import Checkpoint
39
+
40
+ tokenizer, model = Checkpoint("path/to/standalone-model").load("mps")
41
+ ```
42
+
43
+ Copy the **whole standalone model directory** to the target computer. Keep the
44
+ Kev code and its MLX/serving dependencies available. No original model files or
45
+ Hugging Face cache are required on the target computer. This is an MLX format;
46
+ llama.cpp and non-Apple inference have not been implemented or validated.
47
+
48
+ ## Build from the Compressed Experiment
49
+
50
+ The export reads the compressed packed file and its optional `.cloq.safetensors`
51
+ sidecar, plus only the source's small head/tokenizer files and the base's
52
+ `config.json`. The base/PEFT weight files are never read. Rotated and
53
+ activation-order recipes are rejected rather than loaded incorrectly.
54
+
55
+ ```sh
56
+ .venv/bin/python scripts/build_packed_model.py \
57
+ --packed path/to/experimental.safetensors \
58
+ --checkpoint path/to/head-and-tokenizer-directory \
59
+ --config path/to/base-config.json \
60
+ --out path/to/new-standalone-model
61
+ ```
62
+
63
+ The export copies the tokenizer files and retains source hashes in `packed.json`
64
+ and `head.pt`. It does not quantize again, train, or fit a temperature. Native
65
+ 2-bit storage uses more bits than the experiment's five-trits-per-byte storage,
66
+ so the standalone file is larger than the archival packed file. Original
67
+ experimental files and results remain unchanged.
68
+
69
+ ## Verify
70
+
71
+ ```sh
72
+ .venv/bin/python -m pytest tests/test_packed_model.py -q
73
+
74
+ HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-cache \
75
+ .venv/bin/python scripts/profile_packed_model.py \
76
+ --run path/to/standalone-model --mlx-limit-gb 3 --out path/to/memory.json
77
+
78
+ HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-cache KEV_BACKEND=mlx \
79
+ .venv/bin/python -m kev.benchmark --run path/to/standalone-model \
80
+ --suite evals/v7/decision-v7 --device mps --out path/to/new-evaluation
81
+ ```
82
+
83
+ The profiler blocks the original loaders and downloads, and reports both MLX
84
+ allocator memory and sampled process physical footprint. Physical footprint
85
+ includes Metal buffers; RSS alone does not. The MLX allocator cap is not a cap
86
+ on the entire process or the operating system. A test on a larger-memory Mac
87
+ does not establish compatibility with every low-memory device.
88
+
89
+ Native quantized matmul plus a separate compensation adapter uses different
90
+ arithmetic from the original experiment's merged, dequantized bf16 weights.
91
+ Report its newly measured accuracy and Brier score; do not inherit the archival
92
+ 96.7% retention or claim bit-identical probabilities. Context length, option
93
+ count and cached states also consume memory, so report those with memory results.
94
+ `KEV_PREFIX_MAX_TOKENS` controls retained cache size, not input admission.
95
+
96
+ The development partition is already used for recipe selection. Results remain
97
+ development diagnostics, not an untouched final test or a general guarantee.
98
+
99
+ ## Measured 4B Result (2026-10-06)
100
+
101
+ `standalone/kev-4b-T-cloq16-mlx` was evaluated with an empty HF cache in offline
102
+ mode on decision-v7 development: all 1,204 records / 1,468 rows, 1,264 scored
103
+ clean questions, zero rejections and zero truncations. Accuracy is **0.84415**
104
+ (**96.82%** of the archived bf16-MLX accuracy 0.87184), Brier **0.23251**.
105
+ Against the archived packed merged-bf16 probabilities, 1/1,468 argmaxes changed;
106
+ mean per-row max probability difference is 0.00268, maximum 0.06209. The
107
+ original bf16 model was not loaded or re-evaluated for this comparison.
108
+
109
+ On an M4 Pro with 64 GiB RAM, fresh-process probes explicitly blocked the base
110
+ loaders and downloads. MLX 0.32.2, mlx-lm 0.31.3, torch 2.8.0, transformers 5.17.0:
111
+
112
+ | Measurement | Short request | Longer request |
113
+ |---|---:|---:|
114
+ | State tokens / total tokens | 19 / 42 | 3,519 / 3,542 |
115
+ | Load time | 0.57 s | 0.57 s |
116
+ | MLX active weights after load | 1.679 GB | 1.679 GB |
117
+ | Process physical footprint peak during load | 2.396 GB | 2.398 GB |
118
+ | MLX inference peak | 2.013 GB | 3.066 GB |
119
+ | Process physical footprint inference peak | 2.846 GB | 3.862 GB |
120
+ | First inference time | 0.155 s | 6.599 s |
121
+
122
+ These are single-request probes, not a minimum-RAM device certification. No
123
+ 8-GB machine was tested. The 3-GB `mx.set_memory_limit` setting is explicitly a
124
+ **guideline**, not a hard cap; the longer request exceeded it. Use the observed
125
+ physical footprint when judging fit, and leave room for the OS and other apps.
126
+ GB above means decimal GB. The full model folder is approximately 1.709 GB;
127
+ native weights including compensation are 1.679 GB.
128
+
129
+ Artifacts under `runs/ternary-rotation-20261005/standalone-checks/`:
130
+ `kev-4b-T-cloq16-native/report.json`, `native-comparison.json`,
131
+ `memory-short.json`, and `memory-long.json`. The longer synthetic filler probe
132
+ measures memory/latency only and does not establish long-context accuracy.
runtime/kev/__init__.py ADDED
File without changes
runtime/kev/anchors.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Anchor targets: the frozen base model's zero-shot distribution over each training question's options, read from
2
+ next-token letter logits (the same readout as scripts/base_mmlu_probe.py). Used by `kev.train --anchor --anchor_w` as a
3
+ KL target so the fine-tune keeps what the base already knows (PLAN.md: fine-tuning erodes base capability).
4
+
5
+ uv run python -m kev.anchors --base Qwen/Qwen3-4B-Base --suite evals/v6/decision-v6 --out runs/anchors/v6-4b.json
6
+
7
+ Targets are keyed by record id and question id, values keyed by option key, so they align with any option permutation
8
+ the trainer applies; questions whose option set changed (none-of-the-above inserted) are skipped at training time.
9
+ No Jev or other external model is involved: the teacher is the student's own base.
10
+ """
11
+ import argparse
12
+ import json
13
+ from pathlib import Path
14
+
15
+ import torch
16
+ from transformers import AutoModelForCausalLM, AutoTokenizer
17
+
18
+ from kev.api import question_keys
19
+ from kev.device import default_device
20
+ from kev.suite import load_split, write_json
21
+
22
+ LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
23
+
24
+
25
+ def question_prompt(state, q):
26
+ text = state if isinstance(state, str) else " ".join(f"{k}: {v}" for k, v in state.items()) if isinstance(state, dict) else json.dumps(state)
27
+ keys = question_keys(q["type"], q.get("criteria"))
28
+ if q["type"] == "noul": texts = ["No", "Yes"]
29
+ elif q["type"] == "score": texts = list(q["criteria"])
30
+ else: texts = [v if v is not None else k for k, v in q["criteria"].items()]
31
+ prompt = f"{text}\n{q['instructions']}\n" + "\n".join(f"{LETTERS[i]}. {t}" for i, t in enumerate(texts)) + "\nAnswer:"
32
+ return prompt, keys
33
+
34
+
35
+ def build(base, suite, out, split="train", device=None, revision=None, max_options=26, batch=16):
36
+ device = device or default_device()
37
+ tok = AutoTokenizer.from_pretrained(base, revision=revision); tok.padding_side = "left"
38
+ if tok.pad_token is None: tok.pad_token = tok.eos_token
39
+ model = AutoModelForCausalLM.from_pretrained(base, revision=revision, dtype=torch.bfloat16 if device != "cpu" else torch.float32).to(device).eval()
40
+ letter_ids = [tok.encode(" " + L, add_special_tokens=False)[0] for L in LETTERS]
41
+ records = load_split(suite, split)
42
+ jobs = [(r["_meta"]["id"], qid, *question_prompt(r["state"], q)) for r in records for qid, q in r["questions"].items() if len(q.get("criteria", [0, 1])) <= max_options or q["type"] == "noul"]
43
+ targets, skipped = {}, 0
44
+ with torch.no_grad():
45
+ for i in range(0, len(jobs), batch):
46
+ chunk = jobs[i : i + batch]
47
+ enc = tok([p for _, _, p, _ in chunk], return_tensors="pt", padding=True, truncation=True, max_length=3072).to(device)
48
+ logits = model(**enc).logits[:, -1].float()
49
+ for (rid, qid, _, keys), row in zip(chunk, logits):
50
+ if len(keys) > len(letter_ids): skipped += 1; continue
51
+ p = torch.softmax(row[letter_ids[: len(keys)]], -1).tolist()
52
+ targets.setdefault(rid, {})[qid] = dict(zip(keys, p))
53
+ if (i // batch) % 50 == 0: print(f"anchors {i}/{len(jobs)}", flush=True)
54
+ meta = {"base": base, "revision": revision, "suite": str(suite), "split": split, "readout": "zero-shot next-token letter logits",
55
+ "records": len(targets), "questions": sum(len(v) for v in targets.values()), "skipped_over_26_options": skipped}
56
+ write_json(out, {"_meta": meta, "targets": targets})
57
+ print(json.dumps(meta))
58
+ return meta
59
+
60
+
61
+ def main():
62
+ ap = argparse.ArgumentParser()
63
+ ap.add_argument("--base", required=True); ap.add_argument("--suite", required=True); ap.add_argument("--out", required=True)
64
+ ap.add_argument("--split", default="train"); ap.add_argument("--device"); ap.add_argument("--revision")
65
+ a = ap.parse_args()
66
+ if Path(a.out).exists(): raise FileExistsError(a.out)
67
+ build(a.base, a.suite, a.out, a.split, a.device, a.revision)
68
+
69
+
70
+ if __name__ == "__main__":
71
+ main()
runtime/kev/api.py ADDED
@@ -0,0 +1,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """TypeSafe-compatible request/response shapes (POST /v1/systemone) mapped onto the single pointer primitive.
2
+
3
+ Noul -> 2 options [false, true]; answer = p(true)
4
+ Choice -> options 'name' or 'name: desc'; answer = argmax, probabilities by name, confidence
5
+ Score -> options = ordered level descriptions; answer = expected level, legend, probabilities by index
6
+ """
7
+ import json
8
+ import re
9
+ from datetime import datetime
10
+ from typing import Any, Literal, Union
11
+ from pydantic import BaseModel, Field, model_validator
12
+
13
+ JSONContent = Union[str, dict, list, int, float, bool, None]
14
+ MAX_OPTIONS = 255
15
+
16
+
17
+ class Noul(BaseModel):
18
+ type: Literal["noul"]
19
+ instructions: JSONContent = None
20
+ criteria: dict[str, JSONContent] | None = None
21
+
22
+
23
+ class Choice(BaseModel):
24
+ type: Literal["choice"]
25
+ instructions: JSONContent = None
26
+ criteria: dict[str, JSONContent]
27
+
28
+ @model_validator(mode="after")
29
+ def _check(self):
30
+ if not 1 <= len(self.criteria) <= MAX_OPTIONS: raise ValueError(f"criteria must have 1..{MAX_OPTIONS} options")
31
+ return self
32
+
33
+
34
+ class Score(BaseModel):
35
+ type: Literal["score"]
36
+ instructions: JSONContent = None
37
+ criteria: list[JSONContent] = Field(min_length=1, max_length=MAX_OPTIONS)
38
+
39
+
40
+ Question = Union[Noul, Choice, Score]
41
+
42
+
43
+ class SystemOneRequest(BaseModel):
44
+ state: JSONContent
45
+ model: str = "kev-latest"
46
+ questions: dict[str, Question] = Field(min_length=1)
47
+
48
+
49
+ def render(v: JSONContent, indent: int = 0) -> str:
50
+ """Flatten str | object | array into text the model sees. Field names are kept as labels."""
51
+ pad = " " * indent
52
+ if v is None: return ""
53
+ if isinstance(v, (str, int, float, bool)): return str(v)
54
+ if isinstance(v, list): return "\n".join(f"{pad}- {render(x, indent + 1).lstrip()}" for x in v)
55
+ return "\n".join(f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {render(x)}" for k, x in v.items())
56
+
57
+
58
+ def option_text(name: str, desc: JSONContent) -> str:
59
+ return name if desc is None or desc == "" else f"{name}: {render(desc)}"
60
+
61
+
62
+ MONTHS = "January|February|March|April|May|June|July|August|September|October|November|December"
63
+ _DATE = re.compile(rf"\b(?:{MONTHS}) \d{{1,2}}, \d{{4}}\b|\b\d{{4}}-\d{{2}}-\d{{2}}\b")
64
+
65
+
66
+ def date_facts(text: str) -> str:
67
+ """Deterministic date arithmetic for the model: every pair of absolute dates found in `text`, as one sentence each
68
+ ("August 3, 2026 is 12 days after July 22, 2026."). The model cannot subtract dates reliably (issue #8); it can use a
69
+ stated day count. Returns "" when fewer than two dates are found. Dates are listed in order of first appearance."""
70
+ found = []
71
+ for m in _DATE.finditer(text):
72
+ raw = m.group(0)
73
+ try: d = datetime.strptime(raw, "%B %d, %Y") if "," in raw else datetime.strptime(raw, "%Y-%m-%d")
74
+ except ValueError: continue
75
+ if raw not in [r for r, _ in found]: found.append((raw, d))
76
+ facts = []
77
+ for i in range(len(found)):
78
+ for j in range(i + 1, len(found)):
79
+ n = (found[j][1] - found[i][1]).days
80
+ facts.append(f"{found[j][0]} is {abs(n)} day{'s' if abs(n) != 1 else ''} {'after' if n > 0 else 'before'} {found[i][0]}." if n else f"{found[j][0]} is the same day as {found[i][0]}.")
81
+ return " ".join(facts)
82
+
83
+
84
+ def with_date_facts(state):
85
+ """State with a `date_facts` field (object states) or an appended paragraph (string states) when two or more absolute
86
+ dates appear. Opt-in preprocessing (KEV_DATE_FACTS=1 in kev.serve, --date_facts in kev.benchmark)."""
87
+ facts = date_facts(render(state))
88
+ if not facts: return state
89
+ if isinstance(state, dict): return {**state, "date_facts": facts}
90
+ if isinstance(state, list): return state + [{"date_facts": facts}]
91
+ return f"{state}\n\ndate_facts: {facts}"
92
+
93
+
94
+ def question_keys(qtype: str, criteria) -> list[str]:
95
+ """The keys a question's probabilities are reported under, in option order: the criteria names (choice),
96
+ ["false", "true"] (noul), the level indices as strings (score). Labels, targets and anchors use the same keys."""
97
+ if qtype == "choice": return list(criteria)
98
+ if qtype == "noul": return ["false", "true"]
99
+ return [str(i) for i in range(len(criteria))]
100
+
101
+
102
+ def to_record(req: SystemOneRequest):
103
+ """-> internal record for encode(), plus per-question metadata ({"id", "type", "keys", "legend" for score}) to map
104
+ probabilities back."""
105
+ qs, meta = [], []
106
+ for qid, q in req.questions.items():
107
+ m = {"id": qid, "type": q.type, "keys": question_keys(q.type, q.criteria)}
108
+ if q.type == "noul":
109
+ c = q.criteria or {}
110
+ opts = [option_text("no", c.get("false")), option_text("yes", c.get("true"))]
111
+ elif q.type == "choice":
112
+ opts = [option_text(k, v) for k, v in q.criteria.items()]
113
+ else:
114
+ opts = [render(x) for x in q.criteria]
115
+ m["legend"] = dict(zip(m["keys"], opts))
116
+ qs.append({"instr": render(q.instructions), "options": opts, "label": 0}); meta.append(m)
117
+ return {"state": render(req.state), "questions": qs}, meta
118
+
119
+
120
+ # Both confidence formulas mirror TypeSafe's reference adapter, system-one-adapter 0.2.1
121
+ # (src/system_one_adapter/_utils/confidence_metrics.py): p is normalised to sum 1 (all zeros -> uniform), one option -> 1.
122
+ def _normalize(p: list[float]) -> list[float]:
123
+ t = sum(p)
124
+ return [1 / len(p)] * len(p) if t == 0 else [x / t for x in p]
125
+
126
+
127
+ def choice_confidence(p: list[float]) -> float:
128
+ """(p_max - 1/K) / (1 - 1/K): 0 at uniform, 1 at certainty."""
129
+ K = len(p)
130
+ return 1.0 if K == 1 else (max(_normalize(p)) - 1 / K) / (1 - 1 / K)
131
+
132
+
133
+ def score_confidence(p: list[float]) -> float:
134
+ """max(0, 1 - E|level - mode| / D), D = mean absolute deviation of a uniform distribution over the L levels around its
135
+ mean (L-1)/2; mode = first most likely level. 1 when all mass is on one level, 0 at uniform or anything as spread."""
136
+ L = len(p)
137
+ if L == 1: return 1.0
138
+ p = _normalize(p); mode = max(range(L), key=p.__getitem__)
139
+ D = sum(abs(i - (L - 1) / 2) for i in range(L)) / L
140
+ return max(0.0, 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / D)
141
+
142
+
143
+ def round_prob(x: float) -> float:
144
+ """Serialization precision for probabilities and derived scalars. 4 decimals keeps the sum of a rounded distribution
145
+ within TypeSafe's tolerance (|sum - 1| < 0.02) at the 255-option maximum: 255 * 0.00005 < 0.02."""
146
+ return round(float(x), 4)
147
+
148
+
149
+ def to_answers(probs: list[list[float]], meta: list[dict]) -> dict[str, Any]:
150
+ out = {}
151
+ for p, m in zip(probs, meta):
152
+ if m["type"] == "noul":
153
+ out[m["id"]] = {"type": "noul", "noul": round_prob(p[1])}
154
+ elif m["type"] == "choice":
155
+ dist = {k: round_prob(v) for k, v in zip(m["keys"], p)}
156
+ out[m["id"]] = {"type": "choice", "choice": m["keys"][max(range(len(p)), key=lambda i: p[i])], "confidence": round_prob(choice_confidence(p)), "probabilities": dist}
157
+ else:
158
+ score = sum(i * pi for i, pi in enumerate(p))
159
+ out[m["id"]] = {"type": "score", "score": round_prob(score), "legend": m["legend"], "probabilities": {str(i): round_prob(v) for i, v in enumerate(p)}, "confidence": round_prob(score_confidence(p))}
160
+ return out
161
+
162
+
163
+ def output_tokens(tok, answers: dict) -> int:
164
+ """Billing-style figure: tokens of the serialised answers. Not a measure of generation (there is none)."""
165
+ return len(tok(json.dumps(answers), add_special_tokens=False).input_ids)
runtime/kev/autoresearch.py ADDED
@@ -0,0 +1,181 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Unattended research sessions on the round harness, plus the study leaderboard.
2
+
3
+ uv run python -m kev.autoresearch session experiments/rounds/r19.json [...] --spend-start 1377 --spend-cap 300
4
+ uv run python -m kev.autoresearch leaderboard # runs/leaderboard.{jsonl,md} from every trial's result.json
5
+ uv run python -m kev.autoresearch release-check --study X # every seed of a config must pass its gates
6
+ uv run python -m kev.autoresearch compare --studies a,b --reference runs/<study>/<trial>
7
+
8
+ A session runs registered rounds in order through kev.rounds (validate, launch the studies, watch: pull and read each
9
+ finished trial, then the read-out). What it may run is what the specs register: config-only trials through
10
+ kev.experiment.execute_trial with full provenance, and the reads the rule names. It never changes the evaluator, the
11
+ frozen suites or the gates, and it never confirms: test partitions and the locked read are read once, by a deliberate
12
+ `python -m kev.rounds confirm` after a read-out names a candidate (the session prints those commands).
13
+
14
+ Budget: before each round, `modal billing summary` is read; the session stops when metered spend since --spend-start
15
+ plus the round's study budgets (each at least its admission bound, checked by kev.rounds.validate) would pass
16
+ --spend-cap. Each round appends a line to runs/autoresearch-sessions.jsonl.
17
+ """
18
+ import argparse
19
+ import hashlib
20
+ import json
21
+ import subprocess
22
+ import sys
23
+ from collections import defaultdict
24
+ from datetime import datetime, timezone
25
+ from pathlib import Path
26
+
27
+ from kev import rounds
28
+ from kev.experiment import CHOICE_DEFAULTS, DEFAULTS
29
+ from kev.metrics import paired_bootstrap
30
+ from kev.suite import ENCODING, read_json, write_jsonl
31
+
32
+ ROOT = Path(__file__).resolve().parents[1]
33
+ H100_RATE = 3.95
34
+ SNAPSHOT_KEYS = ("snapshot_fractions", "snapshot_every_steps", "snapshot_hub_repo") # when and where a full-weight trial saves intermediate checkpoints: not what it trains
35
+ INFRA_KEYS = ("base", "seed", "base_revision", "dtype", "checkpointing", "batch", "accum", "perm_frac", "shared_prefix", *SNAPSHOT_KEYS) # execution shape, not recipe
36
+ ALL_DEFAULTS = {**DEFAULTS, **CHOICE_DEFAULTS} # what kev.train does when a knob is not given
37
+
38
+
39
+ def config_digest(value):
40
+ """Canonical (key-sorted) digest for config identity; kev.suite.record_digest keeps insertion order for provenance."""
41
+ return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
42
+
43
+
44
+ def knobs(cfg):
45
+ """The recipe a config expresses: every non-default, non-infrastructure parameter."""
46
+ return {k: v for k, v in sorted(cfg.items()) if k not in INFRA_KEYS and v != ALL_DEFAULTS.get(k)}
47
+
48
+
49
+ def recipe(row):
50
+ """A trial's recipe identity: its config without the seed (config_sha256 hashes the seed too) and snapshot schedule."""
51
+ return config_digest({k: v for k, v in row["config"].items() if k != "seed" and k not in SNAPSHOT_KEYS})
52
+
53
+
54
+ # --- sessions --------------------------------------------------------------------------------------------------------
55
+
56
+ def metered_spend():
57
+ out = subprocess.run([sys.executable, "-m", "modal", "billing", "summary", "--json"], capture_output=True, text=True, cwd=ROOT)
58
+ return float(json.loads(out.stdout)["metered_cost"]) if out.returncode == 0 else None
59
+
60
+
61
+ def session(paths, spend_start, spend_cap, spend=metered_spend, log=print):
62
+ """Run each registered round to its read-out while the spend cap allows; returns the ledger entries written."""
63
+ entries = []
64
+ for path in paths:
65
+ spec = rounds.load(path)
66
+ problems = rounds.launchable(spec)
67
+ if problems: log(f"{path}: does not validate, session stops:\n" + "\n".join(problems)); break
68
+ spent, bound = spend(), sum(s["budget"] for s in spec["studies"].values())
69
+ if spent is None: log("cannot read `modal billing summary`; session stops"); break
70
+ if spent - spend_start + bound > spend_cap:
71
+ log(f"round {spec['round']}: spent ${spent - spend_start:.2f} + budgets ${bound:.2f} > cap ${spend_cap:.2f}; session stops"); break
72
+ rounds.launch_studies(spec)
73
+ report = rounds.watch(spec)
74
+ entry = {"round": spec["round"], "spec": str(path), "spend_before": round(spent - spend_start, 2), "budgets": bound, "candidates": report["candidates"],
75
+ "at": datetime.now(timezone.utc).isoformat(timespec="minutes")}
76
+ with (ROOT / "runs/autoresearch-sessions.jsonl").open("a", encoding=ENCODING) as f: f.write(json.dumps(entry) + "\n")
77
+ entries.append(entry)
78
+ for arm in filter(None, report["candidates"].values()):
79
+ for stage in spec.get("confirm", {}):
80
+ log(f"candidate {arm}: uv run python -m kev.rounds launch-reads {path} --stage {stage} --arm {arm} then confirm {path} --stage {stage} --arm {arm}")
81
+ return entries
82
+
83
+
84
+ # --- leaderboard -----------------------------------------------------------------------------------------------------
85
+
86
+ def collect():
87
+ """Every completed trial in runs/*/ with its provenance, as flat leaderboard rows."""
88
+ rows = []
89
+ for result in sorted(ROOT.glob("runs/*/*/result.json")):
90
+ r = read_json(result)
91
+ prov = r.get("provenance", {}); cfg = prov.get("config") or {}
92
+ tr = r.get("transfer") or {}
93
+ rows.append({"study": result.parent.parent.name, "trial": result.parent.name, "legacy": prov.get("legacy_checkpoint", False),
94
+ "base": cfg.get("base"), "seed": cfg.get("seed"), "config": cfg, "config_sha256": prov.get("config_sha256"),
95
+ "suite_sha256": prov.get("suite_sha256"), "transfer_suite_sha256": tr.get("suite_sha256"), "git": prov.get("git_commit", "")[:8],
96
+ "dev_acc": r["clean"]["acc"], "dev_nll": r["clean"]["nll"], "dev_ece": r["clean"]["ece"],
97
+ "transfer_acc": tr.get("clean", {}).get("acc"), "transfer_brier": tr.get("clean", {}).get("brier"),
98
+ "transfer_conf_err": tr.get("clean", {}).get("confident_error_rate"),
99
+ "heldout_pairs": (tr.get("paired_flip") or {}).get("both_correct_rate"),
100
+ "none_present": r["variants"].get("none_present", {}).get("acc"), "perm_flip": r["permutation"]["flip_rate"],
101
+ "gates": r["gates"]["checks"], "gates_passed": r["gates"]["passed"],
102
+ "train_s": (r.get("training_resources") or {}).get("wall_seconds"), "wall_s": r.get("wall_seconds"),
103
+ "est_usd": round(H100_RATE * r.get("wall_seconds", 0) / 3600, 2) if prov.get("device") == "cuda" else None})
104
+ return rows
105
+
106
+
107
+ def leaderboard_md(rows):
108
+ lines = ["# Leaderboard", "", f"Generated {datetime.now(timezone.utc).isoformat(timespec='minutes')} from runs/*/result.json. "
109
+ "Selection on development partitions only; the locked test is never read here. Rounds are decided by their "
110
+ "registered rules (`python -m kev.rounds readout`), not by this table.", "",
111
+ "| study/trial | base | seed | dev acc | transfer acc | Brier | conf-err | held-out pairs | none_present | perm flip | gates | $ | knobs |", "|---|---|---|---|---|---|---|---|---|---|---|---|---|"]
112
+ for r in sorted(rows, key=lambda r: (r["transfer_acc"] is None, -(r["transfer_acc"] or 0))):
113
+ b = (r["base"] or "legacy").split("/")[-1]
114
+ f = lambda x, p=3: "" if x is None else f"{x:.{p}f}"
115
+ lines.append(f"| {r['study']}/{r['trial']} | {b} | {r['seed'] if r['seed'] is not None else ''} | {f(r['dev_acc'])} | {f(r['transfer_acc'])} | {f(r['transfer_brier'])} | {f(r['transfer_conf_err'])} | {f(r['heldout_pairs'],2)} | {f(r['none_present'],2)} | {f(r['perm_flip'],2)} | {'pass' if r['gates_passed'] else 'fail'} | {r['est_usd'] or ''} | {', '.join(f'{k}={v}' for k, v in knobs(r['config']).items()) or 'defaults'} |")
116
+ return "\n".join(lines) + "\n"
117
+
118
+
119
+ def refresh_leaderboard():
120
+ rows = collect()
121
+ write_jsonl(ROOT / "runs/leaderboard.jsonl", rows)
122
+ (ROOT / "runs/leaderboard.md").write_text(leaderboard_md(rows), encoding=ENCODING)
123
+ return rows
124
+
125
+
126
+ def compare(studies, reference, tasks=("mmlu", "paws", "qnli", "emotion", "tweet_offensive", "contrastive_deadline")):
127
+ """Print every trial of the given studies with a record-clustered paired bootstrap on transfer accuracy vs `reference`
128
+ (a runs/<study>/<trial> path). Development-set selection only."""
129
+ def rows(p): return read_json(ROOT / p / "transfer/rows.json")
130
+ ref_rows = rows(reference)
131
+ print(f"reference: {reference}")
132
+ print(f"{'trial':34} {'dev':>6} {'trf':>6} {'brier':>6} {'cerr':>6} {'pairs':>6} {'d trf':>7} {'ci95':>18} knobs / tasks")
133
+ for study in studies:
134
+ for d in sorted((ROOT / "runs" / study).iterdir()):
135
+ if not (d / "result.json").exists(): continue
136
+ r = read_json(d / "result.json"); tr = r.get("transfer")
137
+ if not tr: continue
138
+ cfg = r["provenance"]["config"]
139
+ try:
140
+ b = paired_bootstrap(rows(f"runs/{study}/{d.name}"), ref_rows, metric="acc"); delta, ci = b["macro_acc_delta"], [round(x, 3) for x in b["ci95"]]
141
+ except ValueError:
142
+ delta, ci = float("nan"), "n/a (different suite)"
143
+ print(f"{study + '/' + d.name:34} {r['clean']['acc']:6.3f} {tr['clean']['acc']:6.3f} {tr['clean']['brier']:6.3f} {tr['clean']['confident_error_rate']:6.3f} {tr['paired_flip']['both_correct_rate']:6.2f} {delta:+7.3f} {str(ci):>18} {knobs(cfg)} {({k: round(tr['tasks'][k]['acc'], 2) for k in tasks if k in tr['tasks']})}")
144
+
145
+
146
+ def release_check(study):
147
+ """Release screen across seeds: every trial of the same config in the study must pass its gates (including the 70%
148
+ held-out-pair screen) - a single seed clearing the bar is not enough. Prints the verdict per config."""
149
+ rows = [r for r in collect() if r["study"] == study]
150
+ by_cfg = defaultdict(list)
151
+ for r in rows: by_cfg[recipe(r)].append(r)
152
+ verdicts = {}
153
+ for group in by_cfg.values():
154
+ seeds = sorted(r["seed"] for r in group)
155
+ passed = all(r["gates_passed"] for r in group)
156
+ pairs = [round(r["heldout_pairs"], 2) for r in group]
157
+ failing = sorted({k for r in group for k, v in r["gates"].items() if not v})
158
+ base = group[0]["base"].split("/")[-1]
159
+ verdicts[base] = {"seeds": seeds, "all_gates_passed": passed and len(seeds) >= 2, "heldout_pairs": pairs,
160
+ "transfer_acc": [round(r["transfer_acc"], 3) for r in group], "failing_gates": failing,
161
+ "candidate": passed and len(seeds) >= 2, "trials": [f"{r['study']}/{r['trial']}" for r in group]}
162
+ print(f"{base}: seeds {seeds} pairs {pairs} transfer {verdicts[base]['transfer_acc']} -> {'RELEASE CANDIDATE (gated locked read allowed)' if verdicts[base]['candidate'] else 'not a candidate: ' + ', '.join(failing or ['fewer than two seeds'])}")
163
+ return verdicts
164
+
165
+
166
+ def main():
167
+ ap = argparse.ArgumentParser()
168
+ sub = ap.add_subparsers(dest="cmd", required=True)
169
+ p = sub.add_parser("session"); p.add_argument("specs", nargs="+"); p.add_argument("--spend-start", type=float, required=True); p.add_argument("--spend-cap", type=float, required=True)
170
+ sub.add_parser("leaderboard")
171
+ p = sub.add_parser("release-check"); p.add_argument("--study", required=True)
172
+ p = sub.add_parser("compare"); p.add_argument("--studies", required=True); p.add_argument("--reference", required=True, help="runs/<study>/<trial>")
173
+ a = ap.parse_args()
174
+ if a.cmd == "session": session(a.specs, a.spend_start, a.spend_cap)
175
+ elif a.cmd == "leaderboard": print(f"{len(refresh_leaderboard())} trials -> runs/leaderboard.md")
176
+ elif a.cmd == "release-check": release_check(a.study)
177
+ else: compare(a.studies.split(","), a.reference)
178
+
179
+
180
+ if __name__ == "__main__":
181
+ main()
runtime/kev/benchmark.py ADDED
@@ -0,0 +1,214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Score a predictor on a frozen suite partition (or your own labelled JSONL).
2
+
3
+ uv run python -m kev.benchmark --run runs/<run>/checkpoint --suite evals/<v>/decision-<v> --out runs/<name>
4
+ uv run python -m kev.benchmark --remote http://127.0.0.1:8008 --suite ... --out ... # any System One endpoint
5
+
6
+ Every prediction becomes one row per question (prediction_rows); kev.metrics scores rows; evaluate_records writes
7
+ predictions.jsonl, rows.json and report.json. Predictors live in kev.predictors.
8
+ """
9
+ import argparse
10
+ import functools
11
+ import json
12
+ import math
13
+ import os
14
+ from concurrent.futures import ThreadPoolExecutor
15
+ from pathlib import Path
16
+
17
+ import numpy as np
18
+
19
+ from kev.api import question_keys, with_date_facts
20
+ from kev.checkpoint import LoadOptions
21
+ from kev.contrastive import paired_flip
22
+ from kev.data import api_request, load_records
23
+ from kev.device import default_device
24
+ from kev.metrics import EPSILON, grouped_metrics, metrics, unknowable_report
25
+ from kev.model import ROW_PASS_TOKENS, ContextOverflow
26
+ from kev.predictors import LocalPredictor, RemotePredictor, RotationAveraged
27
+ from kev.suite import CONTEXT, ENCODING, digest, load_split, read_manifest, record_digest, write_json
28
+
29
+
30
+ def labels(q):
31
+ """(option keys, label index) of a labelled request question."""
32
+ keys = question_keys(q["type"], q.get("criteria"))
33
+ return keys, keys.index(q["label"]) if q["type"] == "choice" else int(q["label"])
34
+
35
+
36
+ def validate_distribution(raw, keys):
37
+ if set(raw) != set(keys):
38
+ raise ValueError("probability keys do not match requested options")
39
+ p = np.array([raw[k] for k in keys], dtype=float)
40
+ if not np.isfinite(p).all() or (p < 0).any() or (p > 1).any():
41
+ raise ValueError("non-finite or out-of-range probabilities")
42
+ total = float(p.sum())
43
+ if total <= 0 or abs(total - 1) > max(1e-5, len(keys) * 0.005 + 1e-8):
44
+ raise ValueError(f"invalid probability sum: {total}")
45
+ return p / total, total
46
+
47
+
48
+ def prediction_rows(record, prediction):
49
+ if set(prediction["probabilities"]) != set(record["questions"]):
50
+ raise ValueError("answer IDs differ from request IDs")
51
+ meta = record["_meta"]
52
+ rows = []
53
+ for qid, q in record["questions"].items():
54
+ keys, y = labels(q)
55
+ p, total = validate_distribution(prediction["probabilities"][qid], keys)
56
+ row = {"id": meta["id"], "group": meta["group_id"], "question": qid,
57
+ "source": meta["source"], "task": q["src"], "type": q["type"],
58
+ "variant": meta["variant"], "keys": keys, "label": y, "control_id": meta.get("control_id"),
59
+ "pair_id": meta.get("pair_id"), "sibling": meta.get("sibling"), # suites frozen before parent_id existed stored the parent's id in group_id for variants
60
+ "parent": meta.get("parent_id") or (meta["id"] if meta["variant"] == "clean" else meta["group_id"]),
61
+ "p": p.tolist(), "raw_probability_sum": total, "zero_count": int((p == 0).sum())}
62
+ if "logits" in prediction:
63
+ raw_logits = prediction["logits"][qid]
64
+ if set(raw_logits) != set(keys) or not all(math.isfinite(raw_logits[k]) for k in keys):
65
+ raise ValueError("logit keys or values do not match the requested options")
66
+ row["logits"] = [float(raw_logits[k]) for k in keys]
67
+ row["inference_temperature"] = prediction["inference_temperature"]
68
+ if "kernels" in prediction: row["kernels"] = prediction["kernels"] # a long row off the fp32-exact kernels (LocalPredictor)
69
+ rows.append(row)
70
+ return rows
71
+
72
+
73
+ def summarize(rows, temperature=1.0, heldout_sources=()):
74
+ """Report over benchmark rows. heldout_sources: sources the scored model never trained on; their tasks are also
75
+ reported as a separate block."""
76
+ clean = [r for r in rows if r["variant"] == "clean"]
77
+ tasks = grouped_metrics(clean, "task")
78
+ variants = grouped_metrics(rows, "variant")
79
+ lookup = {(r["id"], r["question"]): r for r in clean}
80
+ diffs, flips = [], []
81
+ for row in rows:
82
+ if row["variant"] == "permuted" and row["type"] == "choice":
83
+ original = lookup[(row["parent"], row["question"])]
84
+ aligned = [row["p"][row["keys"].index(k)] for k in original["keys"]]
85
+ diffs.append(float(np.max(np.abs(np.array(aligned) - original["p"]))))
86
+ flips.append(int(np.argmax(aligned) != np.argmax(original["p"])))
87
+ knowable = [r for r in clean if r["source"] != "unknowable"] # unknowable records are scored on confidence, never on accuracy
88
+ return {"objective": -float(np.mean([v["nll"] for k, v in tasks.items() if not k.startswith("unknowable_") or k.startswith("unknowable_control")])),
89
+ "paired_flip": paired_flip(clean), "unknowable": unknowable_report(clean),
90
+ "clean": metrics(knowable), "tasks": tasks, "variants": variants,
91
+ "heldout_tasks": grouped_metrics([r for r in clean if r["source"] in heldout_sources], "task") if any(r["source"] in heldout_sources for r in clean) else {},
92
+ "permutation": {"n": len(diffs), "mean_max_delta": float(np.mean(diffs)) if diffs else None,
93
+ "flip_rate": float(np.mean(flips)) if flips else None},
94
+ "temperature": temperature, "calibrated_clean": metrics(knowable, temperature),
95
+ "metric_policy": {"version": 2, "selective_ties": "whole_confidence_groups",
96
+ "coverage_at_error": "in-sample maximum over confidence thresholds; not a deployed error guarantee",
97
+ "aurc": "right-step integral over whole confidence groups",
98
+ "confident_error_rate": "high-confidence errors divided by all questions",
99
+ "error_rate_at_0_9": "errors divided by questions accepted at p_max >= 0.9",
100
+ "nll": "exact from logits when recorded; otherwise from floored probabilities",
101
+ "nll_floor": EPSILON, "renormalize_returned_probabilities": True,
102
+ "raw_sums_outside_1e_5": sum(abs(r["raw_probability_sum"] - 1) > 1e-5 for r in rows),
103
+ "returned_zeros": sum(r["zero_count"] for r in rows)}}
104
+
105
+
106
+ def predictions(records, predictor):
107
+ """Yield, in record order, a zero-argument callable that returns predictor(record) or raises what it raised. A
108
+ predictor with `concurrency` > 1 (RemotePredictor: one independent HTTP request per call) keeps that many calls in
109
+ flight on a thread pool; in-process predictors (one GPU) have no `concurrency` and run one record at a time. Either
110
+ way the caller sees each outcome in the same order as the plain sequential loop."""
111
+ workers = getattr(predictor, "concurrency", 1)
112
+ if workers <= 1 or len(records) <= 1:
113
+ for record in records:
114
+ yield functools.partial(predictor, record)
115
+ return
116
+ executor = ThreadPoolExecutor(max_workers=min(workers, len(records)))
117
+ try:
118
+ for future in [executor.submit(predictor, record) for record in records]:
119
+ yield future.result
120
+ finally:
121
+ executor.shutdown(wait=False, cancel_futures=True)
122
+
123
+
124
+ def evaluate_records(records, predictor, directory, temperature=1.0, heldout_sources=(), skip_overlong=False):
125
+ """skip_overlong: for external data that was not admitted to a context (--data, or an eval-only suite frozen as published),
126
+ records the predictor cannot encode are counted in coverage["rejected_records"] and listed in rejected.json instead of
127
+ aborting. Admitted suites never trigger it; reports must state how rejected records enter any headline number."""
128
+ directory = Path(directory)
129
+ directory.mkdir(parents=True, exist_ok=False)
130
+ coverage = {"requested_records": len(records), "requested_questions": sum(len(r["questions"]) for r in records),
131
+ "evaluated_records": 0, "evaluated_questions": 0, "rejected_records": 0, "truncated_records": 0}
132
+ rows, latencies, rejected = [], [], []
133
+ with (directory / "predictions.jsonl").open("w", encoding=ENCODING) as output:
134
+ preds = predictions(records, predictor)
135
+ for record in records:
136
+ try:
137
+ pred = next(preds)()
138
+ new_rows = prediction_rows(record, pred)
139
+ except ContextOverflow as error:
140
+ if skip_overlong:
141
+ coverage["rejected_records"] += 1; rejected.append({"id": record["_meta"]["id"], "error": str(error)}); continue
142
+ coverage["rejected_records"] += 1
143
+ write_json(directory / "failure.json", {"coverage": coverage, "record_id": record["_meta"]["id"], "error_type": type(error).__name__})
144
+ raise
145
+ except Exception as error:
146
+ coverage["rejected_records"] += 1
147
+ write_json(directory / "failure.json", {"coverage": coverage, "record_id": record["_meta"]["id"], "error_type": type(error).__name__})
148
+ raise
149
+ output.write(json.dumps({"request_sha256": record_digest(api_request(record)), "id": record["_meta"]["id"],
150
+ "prediction": pred, "rows": new_rows}, allow_nan=False) + "\n")
151
+ output.flush()
152
+ rows.extend(new_rows)
153
+ latencies.append(pred["latency_ms"])
154
+ coverage["evaluated_records"] += 1
155
+ coverage["evaluated_questions"] += len(new_rows)
156
+ if coverage["evaluated_records"] % 50 == 0:
157
+ print(f"evaluated {coverage['evaluated_records']}/{len(records)}", flush=True)
158
+ write_json(directory / "rows.json", rows)
159
+ if rejected: write_json(directory / "rejected.json", rejected)
160
+ report = summarize(rows, temperature, heldout_sources)
161
+ report.update(coverage=coverage, latency_ms={"median": float(np.median(latencies)), "p95": float(np.quantile(latencies, .95))},
162
+ calibration={"inference_temperature": getattr(predictor, "temperature", None),
163
+ "additional_temperature": temperature, "logits_recorded": all("logits" in r for r in rows)})
164
+ long = [r for r in rows if "kernels" in r]
165
+ if long: # absent when every row ran the exact kernels, so such reports are unchanged
166
+ report["long_rows"] = {"count": len(long), "records": len({r["id"] for r in long}), "kernels": sorted({r["kernels"] for r in long}), "threshold": ROW_PASS_TOKENS}
167
+ write_json(directory / "report.json", report)
168
+ return report, rows
169
+
170
+
171
+ def main():
172
+ ap = argparse.ArgumentParser()
173
+ ap.add_argument("--run", help="checkpoint dir or Hub id (local scoring)")
174
+ ap.add_argument("--remote", help="base URL of a System One-compatible endpoint to score instead of a local checkpoint")
175
+ ap.add_argument("--remote-model", default="kev-latest")
176
+ ap.add_argument("--remote-concurrency", type=int, default=1, help="requests kept in flight against --remote (1 = sequential); rows and their order do not depend on it")
177
+ ap.add_argument("--suite", help="frozen suite directory (scores its development partition)")
178
+ ap.add_argument("--data", help="your own labelled requests, one JSON object per line (kev.data.load_records); an alternative to --suite")
179
+ ap.add_argument("--out", required=True)
180
+ ap.add_argument("--device", choices=["cpu", "mps", "cuda"], default=default_device())
181
+ ap.add_argument("--allow-test", action="store_true")
182
+ ap.add_argument("--split", choices=["development", "calibration", "train"], default="development",
183
+ help="suite partition to score (train: teacher predictions for distillation; --allow-test reads the locked test instead)")
184
+ ap.add_argument("--date_facts", action="store_true", help="apply kev.api.with_date_facts to every state before scoring (the opt-in serving preprocessor); reported in report.json")
185
+ ap.add_argument("--rotations", type=int, default=1, help="average every Choice question over this many cyclic option rotations (kev.predictors.RotationAveraged); 1 = one order")
186
+ a = ap.parse_args()
187
+ if bool(a.run) == bool(a.remote): ap.error("give exactly one of --run or --remote")
188
+ if a.rotations < 1: ap.error("--rotations must be >= 1")
189
+ if a.remote_concurrency < 1: ap.error("--remote-concurrency must be >= 1")
190
+ if bool(a.suite) == bool(a.data): ap.error("give exactly one of --suite or --data")
191
+ if a.data:
192
+ records, heldout, split, source_hash = load_records(a.data), [], "custom", digest(Path(a.data))
193
+ context, skip_overlong = CONTEXT, True
194
+ else:
195
+ split = "test" if a.allow_test else a.split
196
+ records = load_split(a.suite, split, allow_test=a.allow_test)
197
+ manifest = read_manifest(a.suite)
198
+ heldout = manifest["holdout_sources"]; source_hash = digest(Path(a.suite) / "manifest.json")
199
+ context, skip_overlong = manifest.get("context", CONTEXT), bool(manifest.get("eval_only"))
200
+ if a.date_facts:
201
+ records = [{**r, "state": with_date_facts(r["state"])} for r in records]
202
+ predictor = RemotePredictor(a.remote, a.remote_model, os.environ.get("KEV_REMOTE_API_KEY", "local"), concurrency=a.remote_concurrency) if a.remote else LocalPredictor(a.run, a.device, LoadOptions.from_env(), context=context)
203
+ scorer = RotationAveraged(predictor, a.rotations) if a.rotations > 1 else predictor
204
+ report, _ = evaluate_records(records, scorer, a.out, heldout_sources=tuple(heldout), skip_overlong=skip_overlong)
205
+ report.update(suite_sha256=source_hash, data=a.data, date_facts=a.date_facts, rotations=a.rotations, run=a.run or a.remote, split=split,
206
+ calibration_applied=predictor.temperature != 1.0 if not a.remote else None,
207
+ environment=predictor.environment if not a.remote else None, # the kernel set (kev.predictors.kernel_environment)
208
+ remote={"base_url": a.remote, "requested_model": a.remote_model, "served_model": predictor.served_model, "concurrency": a.remote_concurrency} if a.remote else None)
209
+ write_json(Path(a.out) / "report.json", report)
210
+ print(json.dumps({"objective": report["objective"], "clean": report["clean"], "coverage": report["coverage"]}, indent=2))
211
+
212
+
213
+ if __name__ == "__main__":
214
+ main()
runtime/kev/budget.py ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """The resources a Modal trial container asks for and the admission bound on a study's cost: one home for
2
+ modal_app.run_trial/admit_study and kev.rounds.validate. Pure Python (no torch, no modal)."""
3
+ TRIAL_CPU, TRIAL_MEMORY = 4, (65536, 196608) # cores, (request, limit) MiB
4
+ GPU_HOURLY = {"H100": 3.95, "H200": 4.54, "B200": 6.25, "T4": 0.59} # USD per GPU hour
5
+ # Full-weight trials (kev.train --full_ft, kev.full_ft). One GPU keeps the fp32 masters and AdamW moments in host memory:
6
+ # 12 bytes per parameter, ~310 GB for a 27B, plus the staging buffers and the loader; the CPU runs the AdamW step. Several
7
+ # GPUs (FSDP2) keep that state on the GPUs; host memory holds the copy of it a resume point is written from in the
8
+ # background (~307 GB over the ranks) and the gathered checkpoint rank 0 writes (~51 GB; a snapshot being written holds
9
+ # the same, one at a time). Disk: the runs volume stages what a container writes on its local disk, up to two resume
10
+ # points (the new one is complete before the old one goes), the final checkpoint and every snapshot (kept, never deleted),
11
+ # beyond Modal's default 512 GiB quota. A run may plan at most MAX_SNAPSHOTS snapshots (kev.train and
12
+ # kev.experiment.validated_trial refuse more): with a 27B, 2 x 307 + (1 + 8) x 51.2 ≈ 1,075 GB inside the 1 TiB
13
+ # (1,099.5 GB), even if the volume keeps committed files staged. A full disk would fail a snapshot, withhold its resume
14
+ # point and fail the trial (failed.json, not retried), so the cap is a hard limit, not a guideline.
15
+ FULL_FT_SINGLE = 24, (368640, 409600) # 360 / 400 GiB
16
+ FULL_FT_SHARDED = 16, (409600, 471040) # 400 / 460 GiB
17
+ FULL_FT_DISK = 1048576 # MiB of ephemeral disk (1 TiB); Modal bills disk as memory at 20:1
18
+ MAX_SNAPSHOTS = 8 # snapshots one full-weight run may plan (fractions + every-N steps)
19
+ RESUME_POINT_GB, CHECKPOINT_GB = 307, 51.2 # a 27B's resume point (fp32 masters + moments) and bf16 checkpoint, the disk math above
20
+ # A study's limits. Full-weight trials may run to Modal's 24 h cap per attempt and are continued after a timeout (each
21
+ # continuation is a new call that resumes from the trial's last resume point, kev.experiment.continue_trial), so their
22
+ # bound counts every attempt and their budget cap covers one 8 x H200 day (a LoRA study keeps the $250 cap).
23
+ MAX_TIMEOUT = {False: 28800, True: 86400} # a 24 h attempt fits the $1,000 cap on one H200 (3 x 24 h x ~$11/h); on
24
+ # H200:8 (~$41/h) the cap allows ~8 h attempts, a day with the continuations
25
+ MAX_BUDGET = {False: 250, True: 1000}
26
+ # Continuations per full-weight trial: at most 1 + FULL_FT_RETRIES calls, each one attempt. They are spawned by
27
+ # `kev.rounds watch` (or modal_app.py::resume --trial), not by Modal: calls go out with Modal's retries off, because Modal
28
+ # charges a timed-out attempt twice when the container then outlives the 30 s cancellation grace (a timeout and a killed
29
+ # task; a full-weight trial always does, its trainer and volume commits run on), and the retry the kill triggers can start
30
+ # while another attempt of the same input is running. Round 22's trial got two 8 h attempts of the three its bound
31
+ # counted (scripts/modal_retry_probe.py: Retries(2) and killed attempts -> 2 attempts; retries=0 -> exactly 1).
32
+ FULL_FT_RETRIES = 2
33
+ # The attempt lease: a running full-weight attempt holds <study>/<trial>/attempt.json on the kev-leases volume
34
+ # (modal_app.TrialLease), refreshed every LEASE_HEARTBEAT seconds from a thread of its own and marked ended when the
35
+ # attempt exits cleanly. A lease that is not ended and whose heartbeat is younger than LEASE_STALE belongs to a container
36
+ # that may still be writing the trial: a new attempt refuses to start and continue_full_trial waits. LEASE_STALE = 900 s:
37
+ # 30 s Modal's cancellation grace after a timeout, during which the old container keeps running (and committing)
38
+ # + 523 s the longest resume-point commit of the runs volume seen (round 22: 206-523 s); the heartbeat does not wait for it,
39
+ # but a commit the old container started just before its kill is covered even if a heartbeat had waited behind it
40
+ # + 60 s one heartbeat interval (the last heartbeat can be that old when the container dies)
41
+ # = 613 s, rounded up to 15 min for clock skew between the containers and the launcher (seconds, NTP) and a few failed
42
+ # heartbeat commits in a row: a live attempt is mistaken for a dead one only after 15 heartbeats fail in a row.
43
+ # The cost is a continuation that starts up to ~15 min after a killed attempt, against 8 h attempts.
44
+ LEASE_HEARTBEAT, LEASE_STALE = 60, 900
45
+ CPU_HOURLY, MEMORY_GIB_HOURLY = 0.04730, 0.008 # USD per core hour and per GiB hour (memory billed at the limit here)
46
+ # Checkpoint interpolation (scripts/interpolate_checkpoint.py, modal_app.py::interpolate): CPU only. The base backbone stays
47
+ # resident in bf16 (51 GB for Qwen3.8-27B, 25.6B parameters; loading it through transformers also stages the 2.5 GB LM
48
+ # head); the SFT side streams tensor by tensor, one output shard (<= 5 GB) is held until written, and the fp32 arithmetic
49
+ # runs in 16M-element chunks (~200 MB). Peak ~60 GB; the limit leaves room for the loader's transients.
50
+ INTERPOLATE_CPU, INTERPOLATE_MEMORY, INTERPOLATE_TIMEOUT = 8, (81920, 131072), 10800 # cores, 80 / 128 GiB, 3 h
51
+
52
+
53
+ def gpu_count(gpu):
54
+ """Modal's "H200:8" -> 8; a bare type is one GPU."""
55
+ return int(gpu.partition(":")[2] or 1)
56
+
57
+
58
+ def trial_resources(gpu, full_ft=False):
59
+ """(cpu cores, (memory request, limit) MiB) for one trial container."""
60
+ if not full_ft: return TRIAL_CPU, TRIAL_MEMORY
61
+ return FULL_FT_SHARDED if gpu_count(gpu) > 1 else FULL_FT_SINGLE
62
+
63
+
64
+ def trial_disk(full_ft=False):
65
+ """ephemeral_disk (MiB) for one trial container; None keeps Modal's default."""
66
+ return FULL_FT_DISK if full_ft else None
67
+
68
+
69
+ def hourly_rate(gpu, full_ft=False):
70
+ """USD per hour of one trial container of `gpu` ("H200", or "H200:8" for 8 in one container): GPU + CPU + memory at
71
+ the resources trial_resources gives it."""
72
+ kind = gpu.partition(":")[0]
73
+ if kind not in GPU_HOURLY or gpu_count(gpu) < 1: raise ValueError(f"invalid GPU {gpu!r}")
74
+ cpu, memory = trial_resources(gpu, full_ft)
75
+ return GPU_HOURLY[kind] * gpu_count(gpu) + cpu * CPU_HOURLY + (memory[1] + (trial_disk(full_ft) or 0) / 20) / 1024 * MEMORY_GIB_HOURLY
76
+
77
+
78
+ def interpolation_bound(timeout=INTERPOLATE_TIMEOUT, jobs=1):
79
+ """The most `jobs` interpolation containers (INTERPOLATE_CPU, INTERPOLATE_MEMORY; no GPU, no retries) can cost at
80
+ `timeout` seconds each."""
81
+ if timeout <= 0 or jobs < 1: raise ValueError("invalid timeout or job count")
82
+ return (INTERPOLATE_CPU * CPU_HOURLY + INTERPOLATE_MEMORY[1] / 1024 * MEMORY_GIB_HOURLY) * timeout / 3600 * jobs
83
+
84
+
85
+ def compute_bound(gpu, timeout, trials, full_ft=False):
86
+ """The most `trials` containers can cost if each runs to `timeout` seconds, counting a full-weight trial's retries."""
87
+ if timeout <= 0 or trials < 1: raise ValueError("invalid GPU, timeout, or trial count")
88
+ return hourly_rate(gpu, full_ft) * timeout / 3600 * trials * (1 + FULL_FT_RETRIES if full_ft else 1)
89
+
90
+
91
+ def trial_attempts(spawn, label):
92
+ """(attempts used, attempts allowed) for one trial of a spawn record (runs/<study>.spawn.json, modal_app.launch_detached).
93
+ A record written with Modal's retries off ("modal_retries": 0) keeps each trial's attempts in `attempts`, one per entry
94
+ ({"nonce", "call", ...}; a pending entry, written before its spawn, and an abandoned one count too, since their call
95
+ may run), and allows compute_bound's count: 1 + FULL_FT_RETRIES for a full-weight trial, 1 otherwise. An older record
96
+ (Modal retried each call, and nothing local says how many attempts ran) is taken to have used them all."""
97
+ allowed = 1 + FULL_FT_RETRIES if spawn.get("full_ft") else 1
98
+ if spawn.get("modal_retries") != 0: return allowed, allowed
99
+ return max(1, len(spawn.get("attempts", {}).get(label) or [])), allowed
100
+
101
+
102
+ def lease_state(lease, now):
103
+ """'none' (no lease), 'ended' (its attempt exited cleanly), 'fresh' (a heartbeat younger than LEASE_STALE: its container
104
+ may still be writing the trial) or 'stale' (no heartbeat for LEASE_STALE: the container is gone)."""
105
+ if not lease: return "none"
106
+ if lease.get("ended"): return "ended"
107
+ return "fresh" if now - lease["heartbeat"] < LEASE_STALE else "stale"
runtime/kev/calibrate.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Per-workload calibration report from a saved rows.json: what one temperature fitted on *these* rows would do.
2
+
3
+ uv run python -m kev.calibrate --rows runs/kev-9b-wanli-v1/rows.json [--out runs/kev-9b-wanli-v1/calibration.json]
4
+
5
+ A released checkpoint ships one temperature fitted on its in-distribution development rows (scripts/calibrate_checkpoint.py,
6
+ which now refuses such rows unless --allow-in-distribution: new temperatures come from held-out datasets). It
7
+ transfers to our out-of-domain suites, not necessarily to a deployer's workload: Kev-9B on WANLI-256 is served at
8
+ mean confidence 0.82 against accuracy 0.70. This module reports, on one rows file, four arms of the same predictions:
9
+
10
+ raw the checkpoint's logits at T=1 (restored from the recorded inference_temperature)
11
+ shipped as served (the checkpoint's temperature)
12
+ workload one temperature fitted on all these rows (in-sample; the number a deployer would write down)
13
+ workload_oof the same fit, group-disjoint out-of-fold: what fitting on your labelled rows buys on rows you did not fit
14
+
15
+ and a record-clustered paired bootstrap of workload_oof against shipped on ECE, Brier, coverage at <= 5% error and
16
+ AURC. Argmax and accuracy are identical in every arm by construction. Report only: nothing is written to a checkpoint.
17
+ """
18
+ import argparse
19
+ from pathlib import Path
20
+
21
+ from kev.metrics import TEMPERATURE_FIT, fit_temperature, metrics, out_of_fold_rows, paired_bootstrap, raw_row, scored_rows, tempered_row
22
+ from kev.suite import read_json, write_json
23
+
24
+ KEYS = ("n", "acc", "ece", "brier", "nll", "mean_conf", "confident_error_rate", "coverage_at_5pct_error", "aurc")
25
+ COMPARED = ("ece", "brier", "coverage_at_5pct_error", "aurc")
26
+
27
+
28
+ def workload_report(rows, folds=5, seed=0, samples=1000):
29
+ served = scored_rows(rows)
30
+ if not served:
31
+ raise ValueError("no clean, knowable rows to calibrate")
32
+ if any("logits" not in row or "inference_temperature" not in row for row in served):
33
+ raise ValueError("rows without recorded logits and inference_temperature (remote or Jev predictions): a fit would be "
34
+ "relative to the served probabilities, not the checkpoint's raw logits")
35
+ raw = [raw_row(row) for row in served]
36
+ temperature = fit_temperature(raw, **TEMPERATURE_FIT)
37
+ oof, fold_temperatures = out_of_fold_rows(raw, folds, seed, **TEMPERATURE_FIT)
38
+ arms = {"raw": raw, "shipped": served, "workload": [tempered_row(row, temperature) for row in raw], "workload_oof": oof}
39
+ return {"n": len(served), "shipped_temperature": sorted({row.get("inference_temperature", 1.0) for row in served}),
40
+ "workload_temperature": temperature, "fold_temperatures": fold_temperatures,
41
+ "arms": {name: {k: metrics(arm)[k] for k in KEYS} for name, arm in arms.items()},
42
+ "oof_vs_shipped": {metric: paired_bootstrap(oof, served, samples=samples, seed=seed, metric=metric, aggregation="micro") for metric in COMPARED},
43
+ "fit": TEMPERATURE_FIT, "folds": folds, "seed": seed,
44
+ "note": "argmax and accuracy are temperature-invariant; workload is in-sample, workload_oof is the held-out estimate"}
45
+
46
+
47
+ def format_report(report):
48
+ lines = [f"n={report['n']} shipped T={', '.join(f'{t:.2f}' for t in report['shipped_temperature'])} workload T={report['workload_temperature']:.2f}"
49
+ f" folds T=[{', '.join(f'{t:.2f}' for t in report['fold_temperatures'])}]",
50
+ f"{'arm':13}{'acc':>7}{'ece':>7}{'brier':>7}{'nll':>7}{'conf':>7}{'c-err':>7}{'cov@5%':>8}{'aurc':>7}"]
51
+ for name, m in report["arms"].items():
52
+ lines.append(f"{name:13}{m['acc']:7.3f}{m['ece']:7.3f}{m['brier']:7.3f}{m['nll']:7.3f}{m['mean_conf']:7.3f}{m['confident_error_rate']:7.3f}{m['coverage_at_5pct_error']:8.2f}{m['aurc']:7.3f}")
53
+ for metric, b in report["oof_vs_shipped"].items():
54
+ lo, hi = b["ci95"]
55
+ lines.append(f"workload_oof - shipped {metric:24}{b[f'micro_{metric}_delta']:+.3f} 95% CI [{lo:+.3f}, {hi:+.3f}]")
56
+ return "\n".join(lines)
57
+
58
+
59
+ def main():
60
+ ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
61
+ ap.add_argument("--rows", required=True, help="rows.json written by kev.benchmark")
62
+ ap.add_argument("--out", help="where to write the report (default: calibration.json next to the rows)")
63
+ ap.add_argument("--folds", type=int, default=5); ap.add_argument("--seed", type=int, default=0); ap.add_argument("--samples", type=int, default=1000)
64
+ a = ap.parse_args()
65
+ report = workload_report(read_json(a.rows), folds=a.folds, seed=a.seed, samples=a.samples)
66
+ report["rows"] = a.rows
67
+ out = Path(a.out) if a.out else Path(a.rows).with_name("calibration.json")
68
+ write_json(out, report)
69
+ print(format_report(report)); print(f"wrote {out}")
70
+
71
+
72
+ if __name__ == "__main__":
73
+ main()
runtime/kev/checkpoint.py ADDED
@@ -0,0 +1,392 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Trained checkpoints: a run directory or a Hub repo holding a LoRA adapter (or, for a full-weight run, the whole bf16
2
+ backbone), `head.pt` and the tokenizer.
3
+
4
+ Loader rule: `adapter_config.json` present -> a LoRA adapter on `meta.base` at `meta.base_revision`; no adapter and
5
+ `config.json` + `model*.safetensors` (save_pretrained of the backbone, `meta.weights == "full"`) -> the backbone is loaded from
6
+ the checkpoint directory itself, nothing is merged, in the dtype head.pt's `weights_dtype` names (it must match the `dtype`
7
+ save_pretrained wrote to config.json). The tokenizer always comes from the base (both layouts carry a copy).
8
+
9
+ This is the one place that knows the layout of `head.pt` and how a checkpoint becomes a `DecisionModel`:
10
+ `kev.serve`, `kev.benchmark`, `kev.train --init_from`, `kev.publish`, the scripts and the Hugging Face Space all go
11
+ through it. The Space vendors this file next to `model.py` and `api.py` (scripts/publish_space.sh), so it must not
12
+ import the data or suite modules at import time.
13
+
14
+ ck = Checkpoint("jaredpalmer/kev-4b") # or a local run directory; `@tag` pins a Hub revision
15
+ tok, model = ck.load("mps", LoadOptions.from_env())
16
+ ck.meta.temperature # the calibration the checkpoint carries
17
+ """
18
+ import datetime
19
+ import importlib.util
20
+ import json
21
+ import os
22
+ import re
23
+ from dataclasses import dataclass, field
24
+ from pathlib import Path
25
+
26
+ import torch
27
+
28
+ from .model import DecisionModel, is_hybrid, load_tokenizer, pad_id
29
+
30
+ HUB_ID = re.compile(r"[\w.-]+/[\w.-]+(@[\w.-]+)?")
31
+ MTIME_FLOOR = datetime.datetime(2000, 1, 1, tzinfo=datetime.timezone.utc).timestamp() # older file times are zeroed, not dates (Checkpoint.release_date)
32
+
33
+
34
+ def is_hub_id(run):
35
+ return not os.path.isdir(run) and HUB_ID.fullmatch(str(run)) is not None
36
+
37
+
38
+ def resolve_run(run):
39
+ """Local run directory as given, or a Hub repo id like jaredpalmer/kev-4b, optionally pinned to a revision or tag
40
+ with `@` (jaredpalmer/kev-4b@qwen3), downloaded to the HF cache. Returns a str path."""
41
+ if os.path.isdir(run):
42
+ return str(run)
43
+ from huggingface_hub import snapshot_download
44
+ repo, _, revision = str(run).partition("@")
45
+ return snapshot_download(repo, revision=revision or None, allow_patterns=["*.json", "*.safetensors", "*.pt", "*.txt", "*.jinja"])
46
+
47
+
48
+ @dataclass
49
+ class Meta:
50
+ """Contents of `head.pt`. Every reader gets the same defaults for fields older checkpoints did not write.
51
+ `extra` keeps the rest of the file (training args, suite hash, init provenance, temperature fit) so a
52
+ read-modify-write round trip loses nothing."""
53
+ base: str
54
+ head: dict | None = None
55
+ base_revision: str | None = None
56
+ lora: int = 0
57
+ head_dim: int = 256
58
+ option_isolation: bool = False
59
+ special_embeddings: bool = False
60
+ weights_dtype: str = "fp32"
61
+ temperature: float = 1.0
62
+ holdout: list = field(default_factory=list)
63
+ weights: str = "lora" # lora on a base, full backbone, or packed MLX backbone (no base weights)
64
+ extra: dict = field(default_factory=dict)
65
+
66
+ KNOWN = ("base", "head", "base_revision", "lora", "head_dim", "option_isolation", "special_embeddings", "weights_dtype", "temperature", "holdout", "weights")
67
+
68
+ @classmethod
69
+ def from_dict(cls, d):
70
+ return cls(**{k: d[k] for k in cls.KNOWN if k in d}, extra={k: v for k, v in d.items() if k not in cls.KNOWN})
71
+
72
+ def to_dict(self):
73
+ return {**self.extra, **{k: getattr(self, k) for k in self.KNOWN}} # known fields win over a stray key in extra
74
+
75
+
76
+ def read_meta(run):
77
+ return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu"))
78
+
79
+
80
+ def write_meta(run, meta):
81
+ torch.save(meta.to_dict(), f"{run}/head.pt")
82
+
83
+
84
+ @dataclass(frozen=True)
85
+ class LoadOptions:
86
+ """How a checkpoint is turned into a model. Defaults are the exact path every reported number uses; the fields
87
+ are the same knobs the KEV_* environment variables expose to the command-line tools (see from_env).
88
+
89
+ dtype None = fp32, the exact path every reported number uses (bf16 when the checkpoint was trained with a bf16
90
+ backbone). kev.serve defaults to bf16 on CUDA and MPS instead: half the memory, 2-4.5x lower latency on an
91
+ L4 (Kev-4B: 209 -> 118 ms at 101 tokens, 850 -> 189 ms at 330 tokens), probabilities within ~0.01 and
92
+ the same argmax on the checks run so far. KEV_DTYPE=fp32 restores the exact path when serving.
93
+ merge fold the LoRA into the base weights: the delta is computed from the fp32 adapter and added in fp32 with
94
+ one rounding to the load dtype, so a bf16 model holds exactly round(W + delta), the same bits as merging
95
+ an fp32 copy and casting, without the fp32 copy (Kev-9B needed 36 GB of GPU memory to load for that).
96
+ Identical because the Qwen bases are stored in bf16; a base stored in fp32 would be rounded twice.
97
+ Exact in fp32; in bf16 it is faster (~15%) and closer to the fp32 numbers than the unmerged adapter
98
+ (kev-4b, 24 dev records: max |dp| 0.017 vs 0.029, 0 vs 1 argmax flips). Ignored for adapters that carry
99
+ trained token embeddings. A checkpoint trained on a bf16 backbone (--weights_dtype bf16: Kev-27B) keeps
100
+ its adapter unmerged, as it was trained, unless `fused` is asked for (serving); then it is folded the same
101
+ way (Kev-27B served against unmerged: max |dp| 0.009, 0 flips in 280 questions, runs/fused-27b-h200).
102
+ attn attention backend; None = the model default (SDPA on CUDA, eager elsewhere). "sdpa" on MPS measured
103
+ parity with eager and is a few percent faster.
104
+ lora_scale WiSE-FT-style interpolation between base (0) and fine-tuned weights (1), at inference.
105
+ temperature None = the temperature the checkpoint carries (fitted by scripts/calibrate_checkpoint.py); 1.0 = raw logits.
106
+ backend None = torch, the path every reported number uses. "mlx" = kev.mlx_model (Metal kernels for the hybrid
107
+ Qwen3.5 backbones through mlx-lm; the pointer head and encoder are shared; refused for attention-only
108
+ bases, which MPS already runs well). "auto" = mlx when the device is mps, the base is hybrid, mlx-lm is
109
+ installed and fp32 was not asked for (an explicit dtype=float32 means "the exact path"), else torch;
110
+ kev.serve uses auto. The MLX path always merges a LoRA adapter, loads a full-weight checkpoint as
111
+ saved (no merge, no copy: kev.mlx_model.load_full), and ignores `attn` and `dtype` (the backbone runs
112
+ as stored: bf16 for every released checkpoint).
113
+ cuda_graphs replay the serving passes of a hybrid backbone on CUDA (state prefix, question rows on a cached state) as
114
+ CUDA graphs, batched across requests (kev.cuda_graphs, DecisionModel.probs_batch). None = off, the eager
115
+ path every reported number uses; kev.serve turns it on for CUDA. Exact up to floating-point
116
+ reassociation, not bit for bit (the passes are padded to buckets).
117
+ fused rewrite a merged hybrid backbone on CUDA with fused Triton kernels (kev.fused_qwen35; needs
118
+ flash-linear-attention fused_qwen35.FLA_VERSION and refuses any other). None = off; kev.serve turns it on
119
+ for CUDA when fused_available() (KEV_FUSED=0 to decline). Equal to the reference layers up to bf16 rounding.
120
+ """
121
+ dtype: torch.dtype | None = None
122
+ merge: bool = True
123
+ attn: str | None = None
124
+ lora_scale: float = 1.0
125
+ temperature: float | None = None
126
+ backend: str | None = None
127
+ cuda_graphs: bool | None = None
128
+ fused: bool | None = None
129
+
130
+ BACKENDS = (None, "torch", "mlx", "auto")
131
+
132
+ @classmethod
133
+ def from_env(cls, env=os.environ):
134
+ """KEV_DTYPE=bf16|fp16|fp32, KEV_MERGE=0, KEV_ATTN=sdpa|eager, KEV_LORA_SCALE, KEV_TEMPERATURE, KEV_BACKEND=torch|mlx|auto,
135
+ KEV_CUDA_GRAPHS=0|1, KEV_FUSED=0|1.
136
+ For command-line entry points only; library code passes an explicit LoadOptions. Explicit values that equal a
137
+ library default are kept (fp32 as torch.float32, "torch" as a string) so a caller with its own default, like
138
+ kev.serve, can tell "asked for it" from "did not say"."""
139
+ backend = env.get("KEV_BACKEND") or None
140
+ if backend not in cls.BACKENDS: raise ValueError(f"KEV_BACKEND must be one of torch, mlx, auto; got {backend!r}")
141
+ return cls(dtype={"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}.get(env.get("KEV_DTYPE", "")),
142
+ merge=env.get("KEV_MERGE", "1") != "0", attn=env.get("KEV_ATTN") or None,
143
+ lora_scale=float(env.get("KEV_LORA_SCALE", "1")),
144
+ temperature=float(env["KEV_TEMPERATURE"]) if env.get("KEV_TEMPERATURE") else None, backend=backend,
145
+ cuda_graphs={"0": False, "1": True}.get(env.get("KEV_CUDA_GRAPHS", "")),
146
+ fused={"0": False, "1": True}.get(env.get("KEV_FUSED", "")))
147
+
148
+
149
+ def mlx_available():
150
+ try:
151
+ import mlx_lm # noqa: F401
152
+ return True
153
+ except ImportError:
154
+ return False
155
+
156
+
157
+ def fused_available():
158
+ """Whether kev.fused_qwen35 can run: flash-linear-attention importable at its FLA_VERSION (kev.serve's CUDA default;
159
+ not in the serve extra, pinned in the Modal images). An explicit fused=True skips this and fails loudly instead."""
160
+ if importlib.util.find_spec("fla") is None: return False
161
+ try:
162
+ import fla
163
+ from .fused_qwen35 import FLA_VERSION
164
+ except ImportError:
165
+ return False
166
+ return fla.__version__ == FLA_VERSION
167
+
168
+
169
+ class Checkpoint:
170
+ def __init__(self, run):
171
+ self.requested = str(run) # what the caller asked for (a Hub id stays a Hub id in labels)
172
+ self.path = resolve_run(run)
173
+ self.meta = read_meta(self.path)
174
+
175
+ def file(self, name):
176
+ return Path(self.path) / name
177
+
178
+ def adapter_config(self):
179
+ return json.loads(self.file("adapter_config.json").read_text(encoding="utf-8"))
180
+
181
+ @property
182
+ def packed(self):
183
+ found = self.file("packed.json").exists()
184
+ if found != (self.meta.weights == "packed"):
185
+ raise ValueError("packed.json and head.pt weights must agree")
186
+ return found
187
+
188
+ @property
189
+ def full(self):
190
+ """The loader rule (module docstring): True for a full-weight checkpoint, False for a LoRA adapter; head.pt's
191
+ `weights` must agree with the files."""
192
+ if self.packed:
193
+ if self.file("adapter_config.json").exists() or not self.file("config.json").exists() or not self.file("model.safetensors").exists():
194
+ raise ValueError("packed checkpoint requires config.json and model.safetensors, without a PEFT adapter")
195
+ return False
196
+ found = "lora" if self.file("adapter_config.json").exists() else "full" if self.file("config.json").exists() and self.shards() else None
197
+ if found != self.meta.weights:
198
+ raise ValueError(f"{self.path}: head.pt says weights={self.meta.weights!r} but the directory holds "
199
+ f"{ {'lora': 'an adapter', 'full': 'backbone weights'}.get(found, 'neither an adapter nor backbone weights') }")
200
+ return found == "full"
201
+
202
+ def shards(self):
203
+ """The backbone's safetensors files of a full-weight checkpoint (model.safetensors or model-*-of-*.safetensors)."""
204
+ return sorted(Path(self.path).glob("model*.safetensors"))
205
+
206
+ def weights_sha256(self):
207
+ """What a run's provenance pins: the adapter file's sha256, or for full weights the sha256 over every shard's."""
208
+ from .suite import digest # lazy: the Space vendors this module without kev/suite.py
209
+ if not self.full and not self.packed: return digest(self.file("adapter_model.safetensors"))
210
+ import hashlib
211
+ return hashlib.sha256("".join(f"{p.name}:{digest(p)}\n" for p in self.shards()).encode()).hexdigest()
212
+
213
+ def release_date(self):
214
+ """ISO date for the TypeSafe model card: the Hub commit date for a Hub checkpoint. For a local run, or a Hub
215
+ checkpoint offline (its cached files), the UTC date head.pt was written, else the newest of the checkpoint's
216
+ files, else "unknown". An mtime before 2000 is not a date: archives zero it for reproducibility (Kev 1.0's first
217
+ release tarballs did), and it reads as 1970-01-01 (1969-12-31 west of UTC). scripts/build_release_assets.py stamps
218
+ the release date instead."""
219
+ if is_hub_id(self.requested):
220
+ from huggingface_hub import HfApi
221
+ repo, _, revision = self.requested.partition("@")
222
+ try:
223
+ return HfApi().model_info(repo, revision=revision or None).last_modified.date().isoformat()
224
+ except Exception:
225
+ pass
226
+ files = sorted((p.stat().st_mtime for p in Path(self.path).iterdir() if p.is_file()), reverse=True)
227
+ stamp = next((t for t in [self.file("head.pt").stat().st_mtime, *files] if t >= MTIME_FLOOR), None)
228
+ return "unknown" if stamp is None else datetime.datetime.fromtimestamp(stamp, datetime.timezone.utc).date().isoformat()
229
+
230
+ def hybrid_base(self):
231
+ """Whether the backbone has Gated DeltaNet layers (Qwen3.5), read from a config without loading weights: the base's
232
+ for a LoRA adapter, the checkpoint's own config.json for full weights (the backbone that is actually loaded)."""
233
+ from transformers import AutoConfig
234
+ if self.packed or self.full:
235
+ config = AutoConfig.from_pretrained(self.path, local_files_only=self.packed)
236
+ else:
237
+ config = AutoConfig.from_pretrained(self.meta.base, revision=self.meta.base_revision)
238
+ return is_hybrid(config.get_text_config())
239
+
240
+ def backend(self, device, opts=LoadOptions()):
241
+ """The backend `load` will use: LoadOptions.backend resolved ("auto" -> mlx only where it pays and is installed)."""
242
+ if opts.backend not in LoadOptions.BACKENDS: raise ValueError(f"unknown backend {opts.backend!r}")
243
+ if self.packed:
244
+ if opts.backend not in (None, "auto", "mlx") or opts.dtype not in (None, torch.bfloat16):
245
+ raise ValueError("packed checkpoints require the MLX backend with bf16 activations")
246
+ return "mlx"
247
+ if opts.backend != "auto": return opts.backend or "torch"
248
+ exact = opts.dtype is torch.float32 # KEV_DTYPE=fp32: the caller wants the reported-numbers path, not a faster one
249
+ return "mlx" if str(device) == "mps" and not exact and mlx_available() and self.hybrid_base() else "torch"
250
+
251
+ def load(self, device, opts=LoadOptions()):
252
+ """-> (tokenizer, model) in eval mode with the LoRA applied (or the full backbone loaded) and the pointer head loaded. The model is a
253
+ DecisionModel (torch) or an MLXDecisionModel (backend mlx); both expose the same scoring interface."""
254
+ meta = self.meta
255
+ if self.packed:
256
+ self.full # validate the standalone layout before any tensor loading
257
+ self.backend(device, opts)
258
+ if opts.lora_scale != 1 or meta.option_isolation or meta.special_embeddings:
259
+ raise ValueError("packed checkpoints do not support interpolation, isolated options, or trained token embeddings")
260
+ from transformers import AutoTokenizer
261
+ from .packed_model import load_packed_model
262
+ tok = AutoTokenizer.from_pretrained(self.path, local_files_only=True)
263
+ m = load_packed_model(self.path, pad_id(tok), meta.head_dim)
264
+ m.head.load_state_dict(meta.head); m.eval()
265
+ m.head.temperature = meta.temperature if opts.temperature is None else opts.temperature
266
+ return tok, m
267
+ if self.full and opts.lora_scale != 1: raise ValueError("lora_scale interpolates an adapter; a full-weight checkpoint has none")
268
+ tok = load_tokenizer(meta.base, revision=meta.base_revision)
269
+ m = self._load_mlx(tok, opts) if self.backend(device, opts) == "mlx" else self._load_torch(tok, device, opts)
270
+ m.head.load_state_dict(meta.head); m.eval()
271
+ m.head.temperature = meta.temperature if opts.temperature is None else opts.temperature
272
+ return tok, m
273
+
274
+ def _load_mlx(self, tok, opts):
275
+ """-> MLXDecisionModel. A LoRA checkpoint: the base's mlx-lm model with the adapter merged (fp32, one rounding). A
276
+ full-weight checkpoint: mlx-lm's model built from the checkpoint's own config and shards in the dtype head.pt names
277
+ (kev.mlx_model.load_full; strict names/shapes/dtype), nothing merged. Refusals come before the mlx-lm import, which
278
+ would otherwise hide them on a machine without it."""
279
+ full = self.full
280
+ if not full and not opts.merge: raise ValueError("the MLX backend always merges the adapter (KEV_MERGE=0 needs backend=torch)")
281
+ if self.meta.option_isolation: raise ValueError("option_isolation needs the packed mask; not available on the MLX backend")
282
+ if not self.hybrid_base(): raise ValueError(f"the MLX backend is for the hybrid (Qwen3.5) bases; {self.path if full else self.meta.base} is attention-only and runs on MPS with backend=torch")
283
+ dtype = self.saved_dtype() if full else None
284
+ from .mlx_model import MLXDecisionModel, load_base, load_full, merge_lora
285
+ if full:
286
+ lm = load_full(self.path, self.shards(), dtype)
287
+ else:
288
+ lm = load_base(resolve_run(f"{self.meta.base}@{self.meta.base_revision or ''}")) # the base snapshot the torch path already cached
289
+ merge_lora(lm, self.path, opts.lora_scale)
290
+ return MLXDecisionModel(lm, pad_id(tok), head_dim=self.meta.head_dim)
291
+
292
+ def _load_torch(self, tok, device, opts):
293
+ m, merged = self._full_torch(tok, device, opts) if self.full else self._adapted_torch(tok, device, opts)
294
+ serving = str(device).startswith("cuda") and m.hybrid
295
+ if opts.fused and serving and merged: # fused projections need plain (merged or full) weights
296
+ from .fused_qwen35 import fuse
297
+ fuse(m.lm)
298
+ if opts.cuda_graphs and serving:
299
+ from .cuda_graphs import CudaGraphs
300
+ m.graphs = CudaGraphs(m.lm, m.pad_id)
301
+ return m
302
+
303
+ SAVED_DTYPES = {"bf16": "bfloat16", "fp32": "float32"} # head.pt weights_dtype -> the dtype save_pretrained writes to config.json
304
+
305
+ def saved_dtype(self):
306
+ """A full-weight checkpoint's dtype name ("bfloat16" | "float32"): the one head.pt's `weights_dtype` names (bf16 for
307
+ every kev.train --full_ft run: the dtype they were trained in), which must be the dtype save_pretrained recorded in
308
+ config.json; otherwise a mislabelled export would be silently cast (fp32 weights rounded to bf16, or bf16 upcast to
309
+ twice the memory). Both backends load full weights in it."""
310
+ cfg = json.loads(self.file("config.json").read_text(encoding="utf-8"))
311
+ expected, saved = self.SAVED_DTYPES.get(self.meta.weights_dtype), cfg.get("dtype") or cfg.get("torch_dtype")
312
+ if expected is None or saved not in (None, expected):
313
+ raise ValueError(f"{self.path}: config.json records the weights as {saved} but head.pt says weights_dtype={self.meta.weights_dtype!r}")
314
+ return expected
315
+
316
+ def _full_torch(self, tok, device, opts):
317
+ """-> (model, True). Full weights load in `saved_dtype`; an explicit dtype still casts on purpose (fp32: the same
318
+ values computed in fp32). Nothing to merge (`load` refuses lora_scale for full weights on either backend)."""
319
+ meta = self.meta
320
+ return DecisionModel(meta.base, tok, device, head_dim=meta.head_dim, option_isolation=meta.option_isolation,
321
+ dtype=opts.dtype or getattr(torch, self.saved_dtype()), attn=opts.attn, weights=self.path), True
322
+
323
+ def _adapted_torch(self, tok, device, opts):
324
+ """-> (model, whether the adapter was merged): the base with this checkpoint's LoRA."""
325
+ from peft import PeftModel
326
+ meta = self.meta
327
+ dtype, merge = opts.dtype or torch.float32, opts.merge
328
+ if meta.weights_dtype == "bf16":
329
+ # trained with a bf16 backbone (--weights_dtype bf16: Kev-27B, the 35B-A3B MoE whose fused experts need bf16):
330
+ # load it the same way. The exact path keeps the fp32 adapter unmerged; the fused serving path folds it in
331
+ # (one rounding of W + delta, as for every served Kev; parity in runs/serving-27b-*).
332
+ dtype, merge = torch.bfloat16, merge and bool(opts.fused)
333
+ merge = merge and not self.adapter_config().get("trainable_token_indices") # token-trained adapters stay unmerged
334
+ m = DecisionModel(meta.base, tok, device, lora=None, revision=meta.base_revision, head_dim=meta.head_dim,
335
+ option_isolation=meta.option_isolation, dtype=dtype, attn=opts.attn)
336
+ m.lm = PeftModel.from_pretrained(m.lm, self.path, torch_device=str(device)).to(device) # trainable token embeddings, if any, live in the adapter
337
+ if opts.lora_scale != 1:
338
+ for module in m.lm.modules():
339
+ if isinstance(getattr(module, "scaling", None), dict):
340
+ for k in module.scaling: module.scaling[k] *= opts.lora_scale
341
+ m.lora_scale = opts.lora_scale
342
+ if merge: m.lm = m.lm.merge_and_unload() # W += delta: fp32 math, one rounding (see LoadOptions.merge)
343
+ if dtype != torch.float32: m.lm = m.lm.to(dtype)
344
+ return m, merge
345
+
346
+ COMPAT_FIELDS = ("base", "base_revision", "lora", "head_dim", "option_isolation", "special_embeddings", "weights")
347
+
348
+ def warm_start(self, model, ours):
349
+ """Delta training: load this checkpoint's weights (adapter, or full backbone) and pointer head into `model` (a fresh
350
+ DecisionModel built the way `ours` says: with LoRA, or for full-weight training without). `ours` is the Meta the new
351
+ run will save; every architecture field is compared BEFORE loading, because peft and load_state_dict(strict=False)
352
+ load matching keys silently and a half-loaded model still trains and still reports a loss. Returns provenance."""
353
+ from .suite import digest # lazy: the Space vendors this module without kev/suite.py
354
+ for name in self.COMPAT_FIELDS:
355
+ theirs, mine = getattr(self.meta, name), getattr(ours, name)
356
+ if theirs != mine and not (name == "base_revision" and None in (theirs, mine)):
357
+ raise ValueError(f"--init_from {self.path}: {name} is {theirs!r} there and {mine!r} here")
358
+ tensors = self._load_backbone_into(model.lm) if self.full else self._load_adapter_into(model.lm)
359
+ model.head.load_state_dict(self.meta.head)
360
+ return {"init_from": self.requested, "resolved": self.path, "weights_sha256": self.weights_sha256(),
361
+ "head_sha256": digest(self.file("head.pt")), "tensors": tensors}
362
+
363
+ def _load_backbone_into(self, lm):
364
+ """Copy the saved backbone over `lm` shard by shard (the base's copy in memory is replaced, never merged); every
365
+ tensor of `lm` must be covered exactly once. -> tensor count."""
366
+ from safetensors.torch import load_file
367
+ have, seen = set(lm.state_dict()), set()
368
+ for shard in self.shards():
369
+ part = load_file(shard)
370
+ unexpected = sorted(set(part) - have)
371
+ if unexpected: raise ValueError(f"--init_from {self.path}: {shard.name} carries tensors this backbone does not have (e.g. {unexpected[:2]})")
372
+ lm.load_state_dict(part, strict=False); seen |= set(part)
373
+ missing = sorted(have - seen)
374
+ if missing: raise ValueError(f"--init_from {self.path} does not cover {len(missing)} of this backbone's tensors (e.g. {missing[:2]})")
375
+ return len(seen)
376
+
377
+ def _load_adapter_into(self, lm):
378
+ from peft import get_peft_model_state_dict, load_peft_weights, set_peft_model_state_dict
379
+ weights = load_peft_weights(self.path, device="cpu")
380
+ have = set(get_peft_model_state_dict(lm))
381
+ unexpected, missing = sorted(set(weights) - have), sorted(have - set(weights))
382
+ if unexpected:
383
+ raise ValueError(f"--init_from {self.path} carries {len(unexpected)} adapter tensors this model does not have (e.g. {unexpected[:2]}); check --lora_targets / --lora against its adapter_config.json")
384
+ if missing:
385
+ raise ValueError(f"--init_from {self.path} does not cover {len(missing)} of this model's adapter tensors (e.g. {missing[:2]}); check --lora_targets")
386
+ set_peft_model_state_dict(lm, weights)
387
+ return len(weights)
388
+
389
+
390
+ def load(run, device, opts=LoadOptions()):
391
+ """Convenience: Checkpoint(run).load(device, opts)."""
392
+ return Checkpoint(run).load(device, opts)
runtime/kev/compare.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import json
3
+ import math
4
+ from collections import defaultdict
5
+ from pathlib import Path
6
+
7
+ import numpy as np
8
+
9
+ from kev.metrics import paired_bootstrap
10
+ from kev.suite import read_json, write_json
11
+
12
+
13
+ def nll_sensitivity(rows, floor):
14
+ tasks = defaultdict(list)
15
+ for row in rows:
16
+ if row["variant"] == "clean":
17
+ tasks[row["task"]].append(-math.log(max(row["p"][row["label"]], floor)))
18
+ return float(np.mean([np.mean(v) for v in tasks.values()]))
19
+
20
+
21
+ def none_diagnostics(rows):
22
+ out = {}
23
+ for variant in ("none_present", "none_absent"):
24
+ subset = [r for r in rows if r["variant"] == variant]
25
+ values = [r["p"][r["keys"].index("none_of_these")] for r in subset]
26
+ out[variant] = {"n": len(values), "mean_p_none": float(np.mean(values)),
27
+ "p_none_above_half": float(np.mean(np.array(values) > .5))}
28
+ return out
29
+
30
+
31
+ def main():
32
+ ap = argparse.ArgumentParser()
33
+ ap.add_argument("--candidate", required=True)
34
+ ap.add_argument("--reference", required=True)
35
+ ap.add_argument("--out", required=True)
36
+ a = ap.parse_args()
37
+ candidate, reference = Path(a.candidate), Path(a.reference)
38
+ left, right = [read_json(p / "report.json") for p in (candidate, reference)]
39
+ if not left.get("suite_sha256") or left["suite_sha256"] != right.get("suite_sha256"):
40
+ raise ValueError("comparison requires matching frozen-suite hashes")
41
+ lr, rr = [read_json(p / "rows.json") for p in (candidate, reference)]
42
+ result = {"candidate": str(candidate), "reference": str(reference), "suite_sha256": left["suite_sha256"],
43
+ "clean": {"candidate": left["clean"], "reference": right["clean"]},
44
+ "paired": {metric: paired_bootstrap(lr, rr, metric=metric) for metric in ("nll", "acc", "brier")},
45
+ "nll_floor_sensitivity": {str(floor): {"candidate_macro_nll": nll_sensitivity(lr, floor),
46
+ "reference_macro_nll": nll_sensitivity(rr, floor)} for floor in (1e-3, 1e-6, 1e-9)},
47
+ "none_of_the_above": {"candidate": none_diagnostics(lr), "reference": none_diagnostics(rr)},
48
+ "caveats": ["Uncalibrated returned distributions; probabilities normalized for scoring; raw responses preserved.",
49
+ "Historical kev checkpoint trained on all six sources; its heldout-source status does not transfer to this study.",
50
+ "Jev uses a hosted alias and rounds some probabilities to zero; NLL depends on the declared floor.",
51
+ "CPU local latency and hosted network latency are not comparable hardware benchmarks.",
52
+ "This is a development-suite comparison, not the locked final test."]}
53
+ write_json(a.out, result)
54
+ print("task kev acc Jev acc kev NLL Jev NLL")
55
+ for task in left["tasks"]:
56
+ l, r = left["tasks"][task], right["tasks"][task]
57
+ print(f"{task:20} {l['acc']:8.3f} {r['acc']:10.3f} {l['nll']:10.3f} {r['nll']:10.3f}")
58
+ print(json.dumps({"paired": result["paired"], "none_of_the_above": result["none_of_the_above"]}, indent=2))
59
+
60
+
61
+ if __name__ == "__main__":
62
+ main()
runtime/kev/composition.py ADDED
@@ -0,0 +1,321 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import hashlib
3
+ import json
4
+ import random
5
+ from datetime import date, timedelta
6
+
7
+ UNKNOWN = None
8
+ SHAPES = {
9
+ "atom": 0,
10
+ "negation": ("not", 0),
11
+ "conjunction": ("and", 0, 1),
12
+ "disjunction": ("or", 0, 1),
13
+ "exception": ("unless", 0, 1),
14
+ "conditional": ("if", 0, 1, 2),
15
+ "nested_and": ("and", ("and", 0, 1), 2),
16
+ "nested_or": ("or", 0, ("or", 1, 2)),
17
+ "held_and_or": ("and", ("or", 0, 1), 2),
18
+ "held_or_not": ("or", ("and", 0, 1), ("not", 2)),
19
+ "held_conditional": ("if", 0, ("not", 1), 2),
20
+ "final_combination": ("and", ("if", 0, 1, 2), 3),
21
+ "final_negation": ("not", ("or", ("and", 0, 1), 2)),
22
+ "final_exception": ("or", ("unless", 0, 1), 2),
23
+ }
24
+ TRAIN_SHAPES = tuple(list(SHAPES)[:8])
25
+ DEV_SHAPES = tuple(list(SHAPES)[8:11])
26
+ TEST_SHAPES = tuple(list(SHAPES)[11:])
27
+ KINDS = ("lt", "le", "gt", "ge", "eq", "range", "match", "elapsed", "flag")
28
+
29
+
30
+ def atom_value(atom, facts):
31
+ if any(key not in facts for key in atom["fields"]):
32
+ return UNKNOWN
33
+ values = [facts[k] for k in atom["fields"]]
34
+ kind, threshold = atom["kind"], atom["threshold"]
35
+ x = values[0]
36
+ if kind == "lt": return x < threshold
37
+ if kind == "le": return x <= threshold
38
+ if kind == "gt": return x > threshold
39
+ if kind == "ge": return x >= threshold
40
+ if kind == "eq": return x == threshold
41
+ if kind == "range": return threshold <= x <= threshold + 10
42
+ if kind == "match": return x == values[1]
43
+ if kind == "elapsed": return (date.fromisoformat(values[1]) - date.fromisoformat(x)).days <= threshold
44
+ if kind == "flag": return x
45
+ raise ValueError(f"unknown atom {kind}")
46
+
47
+
48
+ def evaluate_rule(tree, atoms, facts):
49
+ if isinstance(tree, int):
50
+ return atom_value(atoms[tree], facts)
51
+ op, *children = tree
52
+ values = [evaluate_rule(t, atoms, facts) for t in children]
53
+ if op == "not": return None if values[0] is None else not values[0]
54
+ if op == "unless":
55
+ a, exception = values
56
+ values = [a, None if exception is None else not exception]
57
+ op = "and"
58
+ if op == "and":
59
+ return False if False in values else None if None in values else True
60
+ if op == "or":
61
+ return True if True in values else None if None in values else False
62
+ if op == "if":
63
+ condition, yes, no = values
64
+ return yes if condition is True else no if condition is False else yes if yes == no else None
65
+ raise ValueError(f"unknown operation {op}")
66
+
67
+
68
+ def atom_text(atom):
69
+ kind, fields, t = atom["kind"], atom["fields"], atom["threshold"]
70
+ key = fields[0]
71
+ words = {"lt": "less than", "le": "at most", "gt": "greater than", "ge": "at least", "eq": "equal to"}
72
+ if kind in words: return f"{key} is {words[kind]} {t}"
73
+ if kind == "range": return f"{key} is between {t} and {t + 10}, including both endpoints"
74
+ if kind == "match": return f"{fields[0]} is the same person as {fields[1]}"
75
+ if kind == "elapsed": return f"the elapsed days from {fields[0]} to {fields[1]} are at most {t}"
76
+ return f"{key} is yes"
77
+
78
+
79
+ def render_rule(tree, atoms, style):
80
+ """Five surface styles for the same rule tree. 0 = logic-like, 1 = 'both/at least one of' lists, 2 = reserved for the
81
+ locked test, 3 = plain prose, 4 = clause per line."""
82
+ if isinstance(tree, int): return atom_text(atoms[tree])
83
+ op, *children = tree
84
+ parts = [render_rule(c, atoms, style) for c in children]
85
+ if style == 3:
86
+ if op == "not": return f"the condition \"{parts[0]}\" fails"
87
+ if op == "and": return f"{parts[0]}, and also {parts[1]}"
88
+ if op == "or": return f"either {parts[0]}, or else {parts[1]}"
89
+ if op == "unless": return f"{parts[0]}, except when {parts[1]}"
90
+ return f"when {parts[0]} the requirement is that {parts[1]}, and when it is not the requirement is that {parts[2]}"
91
+ if style == 4:
92
+ if op == "not": return f"[NOT: {parts[0]}]"
93
+ if op in ("and", "or"): return f"[{'ALL' if op == 'and' else 'ANY'} of: {parts[0]} | {parts[1]}]"
94
+ if op == "unless": return f"[{parts[0]} UNLESS {parts[1]}]"
95
+ return f"[IF {parts[0]} THEN {parts[1]} ELSE {parts[2]}]"
96
+ if op == "not": return f"NOT ({parts[0]})" if style == 0 else f"it is not the case that ({parts[0]})"
97
+ if op in ("and", "or"):
98
+ if style == 0: return f"({parts[0]}) {op.upper()} ({parts[1]})"
99
+ connector = "both" if op == "and" else "at least one of"
100
+ return f"{connector} these conditions hold: [({parts[0]}); ({parts[1]})]"
101
+ if op == "unless": return f"({parts[0]}) holds and the exception ({parts[1]}) does not hold"
102
+ return f"if ({parts[0]}), use ({parts[1]}); otherwise use ({parts[2]})"
103
+
104
+
105
+ # --- random rule trees (PLAN.md: structural diversity, negation anywhere) ---------------------------------------------
106
+
107
+ def skeleton(tree):
108
+ """Structure with leaf identities erased: used to order commutative children independently of atom numbering."""
109
+ if isinstance(tree, int): return "_"
110
+ op, *children = tree
111
+ parts = [skeleton(c) for c in children]
112
+ if op in ("and", "or"): parts = sorted(parts)
113
+ return f"{op}({','.join(parts)})"
114
+
115
+
116
+ def sort_commutative(tree):
117
+ if isinstance(tree, int): return tree
118
+ op, *children = tree
119
+ children = [sort_commutative(c) for c in children]
120
+ if op in ("and", "or"): children = sorted(children, key=skeleton)
121
+ return (op, *children)
122
+
123
+
124
+ def canonical(tree):
125
+ """Structure key: commutative children ordered by skeleton, then leaves renumbered in traversal order, so
126
+ (A and B) or not C and not C or (B and A) share one key."""
127
+ t = relabel(sort_commutative(tree))
128
+ def render(t):
129
+ if isinstance(t, int): return str(t)
130
+ return f"{t[0]}({','.join(render(c) for c in t[1:])})"
131
+ return render(t)
132
+
133
+
134
+ def push_negation(tree):
135
+ """De Morgan normal form, so a random tree equivalent to a held-out shape under negation pushing is also excluded."""
136
+ if isinstance(tree, int): return tree
137
+ op, *children = tree
138
+ if op == "not":
139
+ inner = children[0]
140
+ if isinstance(inner, int): return tree
141
+ iop, *ic = inner
142
+ if iop == "not": return push_negation(ic[0])
143
+ if iop in ("and", "or"): return ("or" if iop == "and" else "and", *[push_negation(("not", c)) for c in ic])
144
+ if iop == "unless": return push_negation(("or", ("not", ic[0]), ic[1]))
145
+ return ("not", push_negation(inner))
146
+ if op == "unless": return ("and", push_negation(children[0]), push_negation(("not", children[1])))
147
+ return (op, *[push_negation(c) for c in children])
148
+
149
+
150
+ def relabel(tree):
151
+ """Renumber leaves in first-appearance order so structure keys do not depend on which atom index was drawn."""
152
+ mapping = {}
153
+ def walk(t):
154
+ if isinstance(t, int):
155
+ mapping.setdefault(t, len(mapping)); return mapping[t]
156
+ return (t[0], *[walk(c) for c in t[1:]])
157
+ return walk(tree)
158
+
159
+
160
+ def structure_keys(tree):
161
+ return {canonical(tree), canonical(push_negation(tree))}
162
+
163
+
164
+ HELD_OUT_KEYS = set().union(*(structure_keys(SHAPES[s]) for s in DEV_SHAPES + TEST_SHAPES))
165
+
166
+
167
+ def random_tree(rng, depth, next_leaf):
168
+ """Random rule over {and, or, not, unless, if} with negation allowed at any position; leaves are fresh atom indices."""
169
+ if depth == 0 or (depth < 3 and rng.random() < 0.25):
170
+ return next_leaf()
171
+ op = rng.choice(["and", "or", "not", "unless", "if", "and", "or"])
172
+ if op == "not": return ("not", random_tree(rng, depth - 1, next_leaf))
173
+ if op == "if": return ("if", random_tree(rng, depth - 1, next_leaf), random_tree(rng, depth - 1, next_leaf), random_tree(rng, depth - 1, next_leaf))
174
+ return (op, random_tree(rng, depth - 1, next_leaf), random_tree(rng, depth - 1, next_leaf))
175
+
176
+
177
+ def sample_trees(n, seed, min_leaves=2, max_leaves=4, exclude=HELD_OUT_KEYS):
178
+ """n distinct rule structures (by canonical key), none structurally equal to a held-out or locked shape, none a
179
+ bare leaf, each with negation somewhere in ~60% of cases. Deterministic in `seed`."""
180
+ rng = random.Random(f"{seed}:trees"); trees, keys = [], set()
181
+ for _ in range(20000):
182
+ if len(trees) >= n: break
183
+ counter = [0]
184
+ def next_leaf():
185
+ counter[0] += 1; return counter[0] - 1
186
+ t = random_tree(rng, rng.choice([2, 2, 3]), next_leaf)
187
+ if isinstance(t, int) or not min_leaves <= counter[0] <= max_leaves: continue
188
+ k = structure_keys(t)
189
+ if k & exclude or k & keys: continue
190
+ keys |= k; trees.append(relabel(t))
191
+ if len(trees) < n: raise ValueError(f"only {len(trees)} distinct structures found")
192
+ return trees
193
+
194
+
195
+ def leaf_indices(tree):
196
+ if isinstance(tree, int): return {tree}
197
+ return set().union(*(leaf_indices(t) for t in tree[1:]))
198
+
199
+
200
+ def make_atoms(tree, rng):
201
+ nouns = rng.sample(["request", "account", "package", "review", "member", "shipment", "entry", "case"], 4)
202
+ atoms = []
203
+ for i in range(max(leaf_indices(tree)) + 1):
204
+ kind = rng.choice(KINDS)
205
+ prefix = nouns[i]
206
+ fields = [f"{prefix} value"]
207
+ if kind == "match": fields = [f"{prefix} signer", f"{prefix} designated approver"]
208
+ elif kind == "elapsed": fields = [f"{prefix} start date", f"{prefix} end date"]
209
+ elif kind == "flag": fields = [f"{prefix} verified"]
210
+ atoms.append({"kind": kind, "fields": fields, "threshold": rng.randint(5, 60)})
211
+ return atoms
212
+
213
+
214
+ def fact_domains(atoms, rng):
215
+ domains = {}
216
+ for a in atoms:
217
+ kind, fs, t = a["kind"], a["fields"], a["threshold"]
218
+ if kind == "match":
219
+ people = rng.sample(["Mira", "Noah", "Aiko", "Ravi", "Sana", "Elin", "Tomas", "Kofi"], 3)
220
+ domains.update({k: people for k in fs})
221
+ elif kind == "elapsed":
222
+ day = date(2027, rng.randint(1, 8), rng.randint(1, 28))
223
+ domains[fs[0]] = [day.isoformat()]
224
+ domains[fs[1]] = [(day + timedelta(days=n)).isoformat() for n in (max(0, t - 1), t, t + 1, t + 10)]
225
+ elif kind == "flag": domains[fs[0]] = [False, True]
226
+ else: domains[fs[0]] = [t - 1, t, t + 1, t + 10, t + 11]
227
+ domains["routing reference"] = [rng.randint(100, 500), rng.randint(501, 999)]
228
+ return domains
229
+
230
+
231
+ def rendered_facts(facts, order):
232
+ def value(v):
233
+ return "yes" if v is True else "no" if v is False else str(v)
234
+ return [f"The {k} is {value(facts[k])}." for k in order]
235
+
236
+
237
+ POLICY_WRAPPERS = {
238
+ 0: "Approve exactly when {rule}. Otherwise deny. The routing reference does not affect eligibility.",
239
+ 1: "Approve exactly when {rule}. Otherwise deny. The routing reference does not affect eligibility.",
240
+ 2: "Approval requires the following rule to be true: {rule}. A false rule means denial. Routing references are irrelevant.",
241
+ 3: "A case is approved when {rule}; any other case is denied. Routing references play no part in the decision.",
242
+ 4: "DECISION RULE {rule} -> approve; otherwise reject. Ignore the routing reference.",
243
+ }
244
+
245
+
246
+ def generate(groups_per_shape, seed, shapes=TRAIN_SHAPES, styles=(0, 1), source="compositional", trees=None):
247
+ """trees: optional {name: tree} to generate from random structures instead of SHAPES (names must not collide)."""
248
+ records = []
249
+ shape_trees = trees if trees is not None else {s: SHAPES[s] for s in shapes}
250
+ for shape, tree in shape_trees.items():
251
+ rng = random.Random(f"{seed}:{shape}")
252
+ for i in range(groups_per_shape):
253
+ atoms = make_atoms(tree, rng)
254
+ domains = fact_domains(atoms, rng)
255
+ for attempt in range(1000):
256
+ facts = {k: rng.choice(v) for k, v in domains.items()}
257
+ label = evaluate_rule(tree, atoms, facts)
258
+ changed = None
259
+ keys = list(domains); rng.shuffle(keys)
260
+ for key in keys:
261
+ for value in domains[key]:
262
+ edited = {**facts, key: value}
263
+ if evaluate_rule(tree, atoms, edited) != label:
264
+ missing = {k: v for k, v in facts.items() if k != key}
265
+ if evaluate_rule(tree, atoms, missing) is None:
266
+ changed = (edited, key)
267
+ break
268
+ if changed: break
269
+ if changed: break
270
+ if changed is None:
271
+ raise ValueError(f"cannot create a decisive edit for {shape}")
272
+ edited, deciding = changed
273
+ nuisance = {**facts, "routing reference": next(v for v in domains["routing reference"] if v != facts["routing reference"])}
274
+ order = list(facts); rng.shuffle(order)
275
+ style = rng.choice(styles)
276
+ rule = render_rule(tree, atoms, style)
277
+ policy = POLICY_WRAPPERS[style].format(rule=rule)
278
+ keys = ["accept", "reject"]; rng.shuffle(keys)
279
+ criteria = {k: "The policy permits this case" if k == "accept" else "The policy does not permit this case" for k in keys}
280
+ group = f"composition/{seed}/{shape}/{i}"
281
+ for kind, a, b in (("relevant", facts, edited), ("irrelevant", facts, nuisance)):
282
+ for sibling, values in (("a", a), ("b", b)):
283
+ result = evaluate_rule(tree, atoms, values)
284
+ state = {"policy": policy, "case": " ".join(rendered_facts(values, order))}
285
+ identifier = f"{group}/{kind}/{sibling}"
286
+ records.append({"state": state, "questions": {"decision": {"type": "choice",
287
+ "instructions": "Apply the policy to this case.", "criteria": dict(criteria),
288
+ "label": "accept" if result else "reject", "src": f"composition_{shape}"}},
289
+ "_meta": {"id": identifier, "group_id": group, "source": source, "variant": "clean",
290
+ "pair_id": f"{group}/{kind}", "sibling": sibling, "pair_kind": kind,
291
+ "family": shape, "family_id": shape, "render_style": style, "structure": canonical(tree),
292
+ "text_sha256": hashlib.sha256(json.dumps(state, sort_keys=True).encode()).hexdigest(),
293
+ "certificate": {"tree": tree, "atoms": atoms, "facts": values, "order": order,
294
+ "deciding_field": deciding, "label": result}}})
295
+ return records
296
+
297
+
298
+ def check_group(records):
299
+ if len(records) != 4: raise ValueError("a composition group requires four records")
300
+ for r in records:
301
+ c = r["_meta"]["certificate"]
302
+ if evaluate_rule(c["tree"], c["atoms"], c["facts"]) != c["label"]:
303
+ raise ValueError("incorrect certificate label")
304
+ if r["state"]["case"] != " ".join(rendered_facts(c["facts"], c["order"])):
305
+ raise ValueError("rendered facts differ from certificate")
306
+ q = r["questions"]["decision"]
307
+ if q["label"] != ("accept" if c["label"] else "reject"):
308
+ raise ValueError("answer differs from certificate")
309
+ for a, b in (records[:2], records[2:]):
310
+ ca, cb = a["_meta"]["certificate"], b["_meta"]["certificate"]
311
+ if ca["order"] != cb["order"] or a["state"]["policy"] != b["state"]["policy"]:
312
+ raise ValueError("pair changes more than one fact")
313
+ if sum(ca["facts"][k] != cb["facts"][k] for k in ca["facts"]) != 1:
314
+ raise ValueError("pair must change exactly one fact")
315
+ expected_flip = a["_meta"]["pair_kind"] == "relevant"
316
+ if (ca["label"] != cb["label"]) != expected_flip:
317
+ raise ValueError("incorrect intervention label")
318
+ missing = {k: v for k, v in ca["facts"].items() if k != ca["deciding_field"]}
319
+ if expected_flip and evaluate_rule(ca["tree"], ca["atoms"], missing) is not None:
320
+ raise ValueError("deciding evidence ablation did not remove the answer")
321
+ return True
runtime/kev/contrastive.py ADDED
@@ -0,0 +1,306 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Programmatic contrastive pairs with certain labels (PLAN.md step 2).
2
+
3
+ Each item is a policy plus a case described by a few sentences. Every sentence carries zero or more explicit facts;
4
+ the label is a pure function `evaluate(facts)` of the facts present, returning UNDETERMINED when a required fact is
5
+ missing. A pair is the same item with exactly one sentence changed so that the label flips.
6
+
7
+ Two code-level checks run on every pair before it is kept (no model anywhere):
8
+ ablation : removing any evidence sentence (one that carries a required fact) makes the label UNDETERMINED,
9
+ so the answer provably lives in the state, not in the instruction, the options, or the policy alone;
10
+ invariance : removing any filler sentence leaves the label unchanged.
11
+ Evidence status is derived from the evaluator, not declared. Both siblings share `pair_id`/`family_id`, and
12
+ splits are assigned by family so siblings never straddle a split.
13
+ """
14
+ import hashlib
15
+ import json
16
+ import random
17
+ from datetime import date, timedelta
18
+
19
+ UNDETERMINED = "UNDETERMINED"
20
+
21
+ NAMES = ["Mira", "Noah", "Priya", "Tomas", "Aiko", "Lena", "Omar", "Sana", "Jonas", "Ravi", "Elin", "Kofi"]
22
+ ROLES = ["account owner", "billing manager", "support agent", "warehouse lead"]
23
+ ITEMS = ["a pair of running shoes", "a desk lamp", "a wireless keyboard", "a rain jacket", "a coffee grinder", "a backpack"]
24
+ PROGRAMS = ["the volunteer driver program", "the apprenticeship", "the rental agreement", "the night-shift roster"]
25
+
26
+
27
+ def _day(d):
28
+ return d.strftime("%B %-d, %Y")
29
+
30
+
31
+ def _need(facts, *keys):
32
+ return all(k in facts for k in keys)
33
+
34
+
35
+ # A family returns (item_a, item_b). An item is {policy, sentences: [(text, facts)], evaluate, question}.
36
+ # The two items must differ in exactly one sentence, and evaluate() must give different labels.
37
+
38
+ def family_return_window(rng):
39
+ window = rng.choice([14, 30, 45, 60]); item = rng.choice(ITEMS); name = rng.choice(NAMES)
40
+ bought = date(2026, rng.randint(1, 9), rng.randint(1, 28))
41
+ def evaluate(f):
42
+ return (f["request"] - f["purchase"]).days <= window if _need(f, "request", "purchase") else UNDETERMINED
43
+ def build(days):
44
+ request = bought + timedelta(days=days)
45
+ return {"policy": f"Returns are accepted only if the return request is submitted within {window} days of the purchase date.",
46
+ "sentences": [(f"{name} bought {item} on {_day(bought)}.", {"purchase": bought}),
47
+ (f"The return request was submitted on {_day(request)}.", {"request": request}),
48
+ (f"The order was paid by card and shipped to {name}'s home address.", {})],
49
+ "evaluate": evaluate, "question": {"type": "noul", "instructions": "Is this return request within the policy window?"}}
50
+ return build(rng.randint(1, window - 1)), build(window + rng.randint(1, 30))
51
+
52
+
53
+ def family_spend_threshold(rng):
54
+ limit = rng.choice([250, 500, 1000, 2500]); name = rng.choice(NAMES); role = rng.choice(ROLES)
55
+ def evaluate(f):
56
+ return ("auto_approved" if f["amount"] <= limit else "director_signoff") if _need(f, "amount") else UNDETERMINED
57
+ def build(amount):
58
+ return {"policy": f"Expense claims of ${limit:,} or less are approved automatically. Claims above ${limit:,} require director sign-off.",
59
+ "sentences": [(f"{name}, the {role}, submitted an expense claim.", {}),
60
+ (f"The claim total is ${amount:,}.", {"amount": amount}),
61
+ ("Receipts were attached for every line item.", {})],
62
+ "evaluate": evaluate, "question": {"type": "choice", "instructions": "How is this claim handled under the policy?",
63
+ "criteria": {"auto_approved": "Approved without further review", "director_signoff": "Requires director sign-off", "rejected": "Rejected outright"}}}
64
+ return build(limit - rng.randint(1, limit // 2)), build(limit + rng.randint(1, limit))
65
+
66
+
67
+ def family_authorization(rng):
68
+ approver, other = rng.sample(NAMES, 2); account = rng.randint(10, 99); amount = rng.choice([40, 120, 350, 900])
69
+ def evaluate(f):
70
+ return f["signer"] == f["approver"] if _need(f, "signer", "approver") else UNDETERMINED
71
+ def build(signer):
72
+ return {"policy": "A refund is authorized only when its sole authorization was signed by someone who may authorize refunds for that account.",
73
+ "sentences": [(f"Only {approver} may authorize refunds for account {account}.", {"approver": approver}),
74
+ (f"The sole authorization for this refund on account {account} was signed by {signer}.", {"signer": signer}),
75
+ (f"The refund amount is ${amount}.", {})],
76
+ "evaluate": evaluate, "question": {"type": "noul", "instructions": "Is the refund authorized?"}}
77
+ return build(approver), build(other)
78
+
79
+
80
+ def family_age_eligibility(rng):
81
+ minimum = rng.choice([16, 18, 21, 25]); name = rng.choice(NAMES); program = rng.choice(PROGRAMS)
82
+ def evaluate(f):
83
+ return f["age"] >= minimum if _need(f, "age") else UNDETERMINED
84
+ def build(age):
85
+ return {"policy": f"Applicants must be at least {minimum} years old to be eligible for {program}.",
86
+ "sentences": [(f"{name} applied to join {program}.", {}),
87
+ (f"{name} is {age} years old.", {"age": age}),
88
+ ("The application form was complete and signed.", {})],
89
+ "evaluate": evaluate, "question": {"type": "noul", "instructions": "Is the applicant eligible?"}}
90
+ return build(minimum + rng.randint(0, 20)), build(minimum - rng.randint(1, 5))
91
+
92
+
93
+ def family_quantity_limit(rng):
94
+ limit = rng.choice([2, 3, 5, 10]); item = rng.choice(ITEMS); name = rng.choice(NAMES)
95
+ def evaluate(f):
96
+ if not _need(f, "qty"): return UNDETERMINED
97
+ return "within_limit" if f["qty"] <= limit else "slightly_over" if f["qty"] <= 2 * limit else "far_over"
98
+ def build(qty):
99
+ return {"policy": f"Customers may order at most {limit} units of any single item per order. Orders up to double the limit are held for review; larger orders are cancelled.",
100
+ "sentences": [(f"{name} placed an order for {item}.", {}),
101
+ (f"The order quantity is {qty}.", {"qty": qty}),
102
+ ("Delivery was requested to a residential address.", {})],
103
+ "evaluate": evaluate, "question": {"type": "choice", "instructions": "What happens to this order?",
104
+ "criteria": {"within_limit": "Processed normally", "slightly_over": "Held for review", "far_over": "Cancelled"}}}
105
+ return build(rng.randint(1, limit)), build(rng.choice([rng.randint(limit + 1, 2 * limit), rng.randint(2 * limit + 1, 4 * limit)]))
106
+
107
+
108
+ def family_deadline(rng):
109
+ name = rng.choice(NAMES); due = date(2026, rng.randint(2, 11), rng.randint(1, 28)); grace = rng.choice([3, 7, 14])
110
+ def evaluate(f):
111
+ if not _need(f, "received", "due"): return UNDETERMINED
112
+ late = (f["received"] - f["due"]).days
113
+ return 0 if late <= 0 else 1 if late <= grace else 2
114
+ def build(offset):
115
+ received = due + timedelta(days=offset)
116
+ return {"policy": f"Reports received by the deadline are on time. Reports received within {grace} days after the deadline are late but accepted. Later reports are refused.",
117
+ "sentences": [(f"The filing deadline for {name}'s report was {_day(due)}.", {"due": due}),
118
+ (f"The report was received on {_day(received)}.", {"received": received}),
119
+ ("The report was submitted through the online portal.", {})],
120
+ "evaluate": evaluate, "question": {"type": "score", "instructions": "How late is this report?", "criteria": ["On time", "Late but accepted", "Refused"]}}
121
+ return build(-rng.randint(0, 10)), build(rng.choice([rng.randint(1, grace), grace + rng.randint(1, 20)]))
122
+
123
+
124
+ # --- ordinal threshold families as Score questions (PLAN.md: the deadline failure is "no idea how ordinal thresholds work").
125
+ # Each has three ordered levels split by two thresholds; one is date-based like deadline but with a different template.
126
+
127
+ def family_warranty_claim(rng):
128
+ name = rng.choice(NAMES); item = rng.choice(ITEMS); bought = date(2026, rng.randint(1, 6), rng.randint(1, 28))
129
+ standard, extended = rng.choice([(90, 365), (180, 730), (365, 1095)]); part = rng.choice(['stitching', 'battery', 'housing', 'zipper', 'switch'])
130
+ def evaluate(f):
131
+ if not _need(f, "claim", "purchase"): return UNDETERMINED
132
+ age = (f["claim"] - f["purchase"]).days
133
+ return 0 if age <= standard else 1 if age <= extended else 2
134
+ def build(age):
135
+ claim = bought + timedelta(days=age)
136
+ return {"policy": f"Warranty claims made within {standard} days of purchase are covered in full. Claims made after that but within {extended} days are covered at half cost. Later claims are not covered.",
137
+ "sentences": [(f"{name} purchased {item} on {_day(bought)}.", {"purchase": bought}),
138
+ (f"A warranty claim for it was filed on {_day(claim)}.", {"claim": claim}),
139
+ (f"The claim describes a defect in the {part}.", {})],
140
+ "evaluate": evaluate, "question": {"type": "score", "instructions": "How is this claim covered?", "criteria": ["Covered in full", "Covered at half cost", "Not covered"]}}
141
+ a = rng.choice([rng.randint(1, standard), rng.randint(standard + 1, extended), extended + rng.randint(1, 200)])
142
+ b = rng.choice([x for x in [rng.randint(1, standard), rng.randint(standard + 1, extended), extended + rng.randint(1, 200)] if evaluate({"claim": bought + timedelta(days=x), "purchase": bought}) != evaluate({"claim": bought + timedelta(days=a), "purchase": bought})] or [a])
143
+ return build(a), build(b)
144
+
145
+
146
+ def family_sla_response(rng):
147
+ name = rng.choice(NAMES); target, breach = rng.choice([(4, 24), (8, 48), (24, 72), (1, 8)])
148
+ unit = "hours"; topic = rng.choice(['a login failure', 'a duplicate charge', 'a missing invoice', 'an export error']); queue = rng.choice(['email', 'chat', 'phone'])
149
+ def evaluate(f):
150
+ if not _need(f, "hours"): return UNDETERMINED
151
+ return 0 if f["hours"] <= target else 1 if f["hours"] <= breach else 2
152
+ def build(h):
153
+ return {"policy": f"Support responses within {target} {unit} meet the service level. Responses after {target} but within {breach} {unit} are a minor breach. Anything slower is a major breach.",
154
+ "sentences": [(f"{name} opened a priority ticket about {topic}.", {}),
155
+ (f"The first response arrived {h} {unit} after the ticket was opened.", {"hours": h}),
156
+ (f"The ticket was routed through the {queue} queue.", {})],
157
+ "evaluate": evaluate, "question": {"type": "score", "instructions": "How does this response time rate against the service level?", "criteria": ["Met", "Minor breach", "Major breach"]}}
158
+ levels = [rng.randint(1, target), rng.randint(target + 1, breach), breach + rng.randint(1, 100)]
159
+ a, b = rng.sample(levels, 2)
160
+ return build(a), build(b)
161
+
162
+
163
+ def family_late_fee(rng):
164
+ name = rng.choice(NAMES); grace, cap = rng.choice([(5, 30), (10, 60), (15, 45)]); amount = rng.choice([120, 450, 980, 2300]); method = rng.choice(['bank transfer', 'card', 'cheque'])
165
+ def evaluate(f):
166
+ if not _need(f, "days_late"): return UNDETERMINED
167
+ return 0 if f["days_late"] <= grace else 1 if f["days_late"] <= cap else 2
168
+ def build(d):
169
+ return {"policy": f"Invoices paid within {grace} days after the due date incur no fee. Payments between {grace + 1} and {cap} days late incur a 2% fee. Payments later than {cap} days incur a 10% fee and a hold on the account.",
170
+ "sentences": [(f"{name}'s invoice for ${amount:,} fell due last quarter.", {}),
171
+ (f"Payment was received {d} days after the due date.", {"days_late": d}),
172
+ (f"The payment was made by {method}.", {})],
173
+ "evaluate": evaluate, "question": {"type": "score", "instructions": "Which fee tier applies?", "criteria": ["No fee", "2% fee", "10% fee and account hold"]}}
174
+ levels = [rng.randint(0, grace), rng.randint(grace + 1, cap), cap + rng.randint(1, 60)]
175
+ a, b = rng.sample(levels, 2)
176
+ return build(a), build(b)
177
+
178
+
179
+ def family_volume_discount(rng):
180
+ name = rng.choice(NAMES); item = rng.choice(ITEMS); t1, t2 = rng.choice([(10, 50), (25, 100), (5, 20), (100, 500)]); dest = rng.choice(['warehouse', 'storefront', 'branch office'])
181
+ def evaluate(f):
182
+ if not _need(f, "units"): return UNDETERMINED
183
+ return 0 if f["units"] < t1 else 1 if f["units"] < t2 else 2
184
+ def build(u):
185
+ return {"policy": f"Orders of fewer than {t1} units are charged the list price. Orders of {t1} to {t2 - 1} units receive the volume discount. Orders of {t2} units or more receive the wholesale rate.",
186
+ "sentences": [(f"{name} placed a business order for {item}.", {}),
187
+ (f"The order is for {u} units.", {"units": u}),
188
+ (f"Delivery is to a {dest}.", {})],
189
+ "evaluate": evaluate, "question": {"type": "score", "instructions": "Which pricing tier applies?", "criteria": ["List price", "Volume discount", "Wholesale rate"]}}
190
+ levels = [rng.randint(1, t1 - 1), rng.randint(t1, t2 - 1), t2 + rng.randint(0, t2)]
191
+ a, b = rng.sample(levels, 2)
192
+ return build(a), build(b)
193
+
194
+
195
+ def family_shipping_delay(rng):
196
+ """Day-precision date arithmetic with small thresholds (the skill `deadline` tests): promised vs delivered dates."""
197
+ name = rng.choice(NAMES); item = rng.choice(ITEMS); promised = date(2026, rng.randint(1, 11), rng.randint(1, 28))
198
+ minor, major = rng.choice([(1, 4), (2, 7), (3, 10)]); carrier = rng.choice(["the courier", "the postal service", "a freight partner"])
199
+ def evaluate(f):
200
+ if not _need(f, "delivered", "promised"): return UNDETERMINED
201
+ late = (f["delivered"] - f["promised"]).days
202
+ return 0 if late <= minor else 1 if late <= major else 2
203
+ def build(offset):
204
+ delivered = promised + timedelta(days=offset)
205
+ return {"policy": f"Deliveries up to {minor} day{'s' if minor > 1 else ''} after the promised date count as on time. Deliveries {minor + 1} to {major} days after it are a minor delay and earn a shipping refund. Later deliveries are a major delay and earn a full refund.",
206
+ "sentences": [(f"{name} ordered {item} with delivery promised for {_day(promised)}.", {"promised": promised}),
207
+ (f"The parcel was delivered on {_day(delivered)}.", {"delivered": delivered}),
208
+ (f"It was shipped by {carrier}.", {})],
209
+ "evaluate": evaluate, "question": {"type": "score", "instructions": "How is this delivery classified?", "criteria": ["On time", "Minor delay: shipping refund", "Major delay: full refund"]}}
210
+ levels = [rng.randint(-3, minor), rng.randint(minor + 1, major), major + rng.randint(1, 30)]
211
+ a, b = rng.sample(levels, 2)
212
+ return build(a), build(b)
213
+
214
+
215
+ FAMILIES = {"return_window": family_return_window, "spend_threshold": family_spend_threshold, "authorization": family_authorization,
216
+ "age_eligibility": family_age_eligibility, "quantity_limit": family_quantity_limit, "deadline": family_deadline,
217
+ "warranty_claim": family_warranty_claim, "sla_response": family_sla_response, "late_fee": family_late_fee, "volume_discount": family_volume_discount,
218
+ "shipping_delay": family_shipping_delay}
219
+ ORDINAL_FAMILIES = ("warranty_claim", "sla_response", "late_fee", "volume_discount", "shipping_delay") # trainable Score-threshold families (deadline stays held out)
220
+
221
+
222
+ def label_of(item, drop=None):
223
+ facts = {}
224
+ for i, (_, f) in enumerate(item["sentences"]):
225
+ if i != drop: facts.update(f)
226
+ return item["evaluate"](facts)
227
+
228
+
229
+ def check_pair(a, b):
230
+ """None if the pair is valid, else the reason it is rejected."""
231
+ la, lb = label_of(a), label_of(b)
232
+ if UNDETERMINED in (la, lb): return "label_undetermined"
233
+ if la == lb: return "labels_equal"
234
+ if sum(x[0] != y[0] for x, y in zip(a["sentences"], b["sentences"])) != 1 or len(a["sentences"]) != len(b["sentences"]):
235
+ return "not_exactly_one_sentence_differs"
236
+ for item, label in ((a, la), (b, lb)):
237
+ evidence = 0
238
+ for i, (_, facts) in enumerate(item["sentences"]):
239
+ got = label_of(item, drop=i)
240
+ if facts and got != UNDETERMINED: return "ablation_failed" # evidence removed must make it undeterminable
241
+ if not facts and got != label: return "invariance_failed" # filler removed must not change the label
242
+ evidence += bool(facts)
243
+ if evidence < 1: return "no_evidence_sentence"
244
+ return None
245
+
246
+
247
+ def to_request(item, family, pair_id, sibling, rng):
248
+ order = list(range(len(item["sentences"]))); rng.shuffle(order)
249
+ sentences = [item["sentences"][i][0] for i in order]
250
+ q = {**item["question"], "label": label_of(item), "src": f"contrastive_{family}"}
251
+ return {"state": {"policy": item["policy"], "case": " ".join(sentences)}, "questions": {"decision": q},
252
+ "_meta": {"source": "contrastive", "family": family, "family_id": f"{family}/{pair_id}", "pair_id": pair_id, "sibling": sibling,
253
+ "repo": None, "revision": None, "split": "generated", "row": pair_id, "id": f"contrastive/{family}/{pair_id}/{sibling}",
254
+ "text_sha256": hashlib.sha256(" ".join(sentences).casefold().encode()).hexdigest(),
255
+ "row_sha256": hashlib.sha256(json.dumps([t for t, _ in item["sentences"]]).encode()).hexdigest()}}
256
+
257
+
258
+ def generate(n_pairs_per_family, seed, families=None):
259
+ """Sibling-adjacent records plus a per-family report of rejections."""
260
+ records, report = [], {}
261
+ for family in (families or FAMILIES):
262
+ rng = random.Random(f"{seed}:{family}")
263
+ kept, rejected, reasons, attempts = 0, 0, {}, 0
264
+ while kept < n_pairs_per_family and attempts < 50 * n_pairs_per_family:
265
+ attempts += 1
266
+ a, b = FAMILIES[family](rng)
267
+ why = check_pair(a, b)
268
+ if why:
269
+ rejected += 1; reasons[why] = reasons.get(why, 0) + 1; continue
270
+ pair_id = f"{seed}-{family}-{kept:04d}"
271
+ order_seed = rng.getrandbits(64)
272
+ records += [to_request(a, family, pair_id, "a", random.Random(order_seed)),
273
+ to_request(b, family, pair_id, "b", random.Random(order_seed))]
274
+ kept += 1
275
+ if kept < n_pairs_per_family:
276
+ raise ValueError(f"{family}: only {kept}/{n_pairs_per_family} pairs passed checks ({reasons})")
277
+ report[family] = {"pairs": kept, "rejected": rejected, "reasons": reasons}
278
+ return records, report
279
+
280
+
281
+ def paired_flip(rows):
282
+ """Pair-level metrics from benchmark rows carrying pair_id/sibling. A model that ignores the state cannot flip."""
283
+ by_pair = {}
284
+ for r in rows:
285
+ if r.get("pair_id"):
286
+ key = (r["pair_id"], r.get("question", "decision"))
287
+ pair = by_pair.setdefault(key, {})
288
+ if r["sibling"] in pair:
289
+ raise ValueError("duplicate contrastive sibling")
290
+ pair[r["sibling"]] = r
291
+ if not by_pair:
292
+ return None
293
+ if any(set(p) != {"a", "b"} for p in by_pair.values()):
294
+ raise ValueError("incomplete contrastive pair")
295
+ prediction = lambda r: r["keys"][max(range(len(r["p"])), key=r["p"].__getitem__)]
296
+ truth = lambda r: r["keys"][r["label"]]
297
+ relevant = [p for p in by_pair.values() if truth(p["a"]) != truth(p["b"])]
298
+ invariant = [p for p in by_pair.values() if truth(p["a"]) == truth(p["b"])]
299
+ both = lambda ps: sum(all(prediction(r) == truth(r) for r in p.values()) for p in ps) / len(ps) if ps else None
300
+ result = {"pairs": len(relevant),
301
+ "flip_rate": sum(prediction(p["a"]) != prediction(p["b"]) for p in relevant) / len(relevant) if relevant else None,
302
+ "both_correct_rate": both(relevant)}
303
+ if invariant:
304
+ result.update(invariant_pairs=len(invariant), invariance_rate=sum(prediction(p["a"]) == prediction(p["b"]) for p in invariant) / len(invariant),
305
+ invariant_both_correct_rate=both(invariant))
306
+ return result
runtime/kev/cuda_graphs.py ADDED
@@ -0,0 +1,375 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """CUDA graphs for serving the hybrid (Qwen3.5) backbones on CUDA, batched across requests.
2
+
3
+ A served request is two kinds of forward pass: the state pass (the document, on a prefix-cache miss) and the row pass (one
4
+ causal row per question, continuing the cached state). At a few hundred tokens the GPU finishes a pass long before Python
5
+ finishes issuing it: a pass launches ~2,000 kernels, and the flash-linear-attention wrapper around every DeltaNet layer
6
+ costs ~1 ms of CPU. Replaying a captured graph issues the same kernels in one call. And because one request leaves most of
7
+ the GPU idle, kev.serve hands this module every request that is waiting when the GPU frees up: one state pass computes all
8
+ their new states, one row pass all their questions.
9
+
10
+ Two sets of buffers, shared by every graph:
11
+ - the state bank: per layer the attention keys/values [GRAPH_STATES, heads, BANK_WIDTH, dim], each state right-aligned at
12
+ the end, and the DeltaNet conv/recurrent states. The state pass writes new states into it; cached states are copied in.
13
+ One fixed layout, so a row pass reads any state the same way whatever state pass wrote it.
14
+ - the row buffers: [rows, heads, Sr + Lb, dim] per attention layer and the DeltaNet states per question row. The row pass
15
+ starts by gathering each row's state from the bank (one index_select per tensor, inside the graph), then writes the
16
+ row's own keys after it.
17
+ Both are flat buffers viewed at the graph's shape, so memory is fixed however many shapes appear. Graph keys are kept
18
+ coarse because a batch's shape changes with the traffic: state passes by (states, state bucket), row passes by (rows,
19
+ row bucket, Sr = the longest state rounded up to a power of two, which only lengthens attention).
20
+
21
+ A graph has fixed shapes, so passes are padded to buckets: token counts to bucket() (under a quarter more), row and state counts
22
+ to count_bucket() (empty entries are masked). The padding is masked exactly:
23
+ - the state pass is LEFT-padded. The DeltaNet layers see zeroed inputs at the pads (their padding mask), so keys, values
24
+ and queries are zero and the recurrent state is still zero when the real tokens start; the causal convolution sees
25
+ the same zeros its own padding supplies. The attention layers mask the pads as keys, and a pad query attends to itself
26
+ so no row is fully masked (a NaN there would survive the zeroing: NaN * 0 = NaN). Real tokens keep their position ids.
27
+ - question rows are right-padded (after every real token, as the eager path already does); a state's keys are
28
+ right-aligned in the bucket and the slots before them are masked.
29
+ So the results equal the eager passes up to floating-point reassociation (another chunking of the DeltaNet scan, other
30
+ GEMM shapes), not bit for bit, and a request's answers do not depend on what else shares its batch.
31
+
32
+ Capturing a graph costs ~0.4 s (an H100, Kev-4B), so the first pass of a new bucket runs the same code eagerly from the
33
+ same buffers and the bucket joins `pending`; capture_pending() captures them later (kev.serve does it when no request is
34
+ waiting). Replays run one at a time (kev.serve has one model thread), each refills what it reads, and its results are
35
+ copied out before the next replay; that is what makes the shared buffers and the shared memory pool safe.
36
+ """
37
+ from collections import OrderedDict
38
+ from typing import NamedTuple
39
+
40
+ import torch
41
+ from transformers import DynamicCache
42
+ from transformers.cache_utils import DynamicLayer, LinearAttentionLayer
43
+
44
+ # Limits of the graphed passes (not the model's context: that is kev.model.MAX_STATE / SERVE_MAX_STATE)
45
+ GRAPH_TOKENS = 32768 # rows x positions one row pass may hold in the attention buffers (about 1 GB on Kev-4B and 9B)
46
+ GRAPH_STATE = 1024 # longest state bucket the state pass graphs. A longer state pass is compute-bound, and the padding plus
47
+ # the explicit mask (no causal flash attention) made a 2,200-token one slower as a graph (L40S, Kev-4B:
48
+ # 223 vs 208 ms per request), so it runs eagerly
49
+ GRAPH_ROW = 1024 # longest question-row bucket the row pass graphs
50
+ GRAPH_ROWS = 32 # question rows per graphed pass; more run as several replays
51
+ GRAPH_STATES = 16 # states per state pass: the bank's entries
52
+ BANK_WIDTH = 4096 # positions per bank entry (states are right-aligned in it; about 2 GB for 16 entries on Kev-4B and 9B)
53
+ GRAPHS_KEPT = 256 # captured graphs kept, least recently used evicted (a busy server met ~160 on mixed traffic)
54
+ PASS_TOKENS = 256 # tokens a pass costs however few it holds: below about this many, a pass is bound by reading the
55
+ # weights, not by compute (bf16 GEMMs: peak FLOP/s over memory bandwidth is ~200-400 on H100, H200, B200, L40S)
56
+ HOT_BUCKET = 3 # eager passes after which a busy server captures a bucket's graph anyway (capture_due)
57
+
58
+
59
+ def bucket(n, steps=8, floor=16):
60
+ """n rounded up to one of `steps` steps per power of two, in steps of at least `floor`: the padding adds less than
61
+ 2/steps of n (just past a power of two, one step is almost that), or less than `floor` tokens for short n."""
62
+ step = max(floor, (1 << (n - 1).bit_length()) // steps)
63
+ return -(-n // step) * step
64
+
65
+
66
+ def pow2(n):
67
+ return 1 << (n - 1).bit_length()
68
+
69
+
70
+ def count_bucket(n):
71
+ """A row or state count rounded up to 1, 2, 3, 4, 6, 8, 12, 16, 24, 32: the empty entries padding a pass add less than
72
+ half, and a traffic mix meets few distinct counts."""
73
+ return bucket(n, steps=4, floor=1)
74
+
75
+
76
+ def length_groups(lengths, cap):
77
+ """Indices grouped into padded passes of at most `cap` items, in order of length, computing the fewest tokens: a pass
78
+ costs its padded size, count_bucket(items) x bucket(longest), but at least PASS_TOKENS (below that, reading the
79
+ weights dominates). Optimal among groupings of consecutive lengths (dynamic programming; ties keep larger passes)."""
80
+ order = sorted(range(len(lengths)), key=lengths.__getitem__)
81
+ padded = [bucket(lengths[i]) for i in order]
82
+ counts = [0] + [count_bucket(n) for n in range(1, cap + 1)]
83
+ best, cut = [0] + [float("inf")] * len(order), [0] * (len(order) + 1)
84
+ for j in range(1, len(order) + 1):
85
+ for i in range(max(0, j - cap), j):
86
+ cost = best[i] + max(PASS_TOKENS, counts[j - i] * padded[j - 1])
87
+ if cost < best[j]: best[j], cut[j] = cost, i
88
+ groups, j = [], len(order)
89
+ while j: groups.append(order[cut[j]:j]); j = cut[j]
90
+ return groups[::-1]
91
+
92
+
93
+ class Request(NamedTuple):
94
+ """One request for CudaGraphs.run."""
95
+ state: list # the state's token ids
96
+ positions: list # and their position ids
97
+ rows: list # its question rows, [(token ids, position ids)]
98
+ cached: object # the cached state (a DynamicCache) it continues, or None for a new state
99
+ keep: bool # return a new state's cache (to keep in the prefix cache)
100
+ picks: list # per row, the positions whose hidden states are wanted
101
+
102
+
103
+ class _Row(NamedTuple):
104
+ """One question row of a batch: its tokens, its state's bank entry and length, and where its picks go in the output."""
105
+ ids: list
106
+ pos: list
107
+ entry: int
108
+ state_len: int
109
+ picks: list
110
+ at: int
111
+
112
+
113
+ class BufferKV(DynamicLayer):
114
+ """An attention cache layer over preallocated [N, heads, T, dim] views: the first `filled` positions hold the cached
115
+ state, a pass writes its keys/values after them instead of concatenating (no allocation, no Python state change,
116
+ so one layer object serves the warm-up, the capture and every replay)."""
117
+
118
+ def __init__(self, keys, values, filled):
119
+ super().__init__()
120
+ self.buffers, self.filled, self.is_initialized = (keys, values), filled, True
121
+ self.keys, self.values = keys[..., :filled, :], values[..., :filled, :]
122
+
123
+ def update(self, key_states, value_states, *args, **kwargs):
124
+ end = self.filled + key_states.shape[-2]
125
+ for buf, new in zip(self.buffers, (key_states, value_states)): buf[..., self.filled:end, :].copy_(new)
126
+ return tuple(buf[..., :end, :] for buf in self.buffers)
127
+
128
+
129
+ def set_linear(layer, conv, recurrent, previous):
130
+ """Point a DeltaNet cache layer at these conv/recurrent state tensors (transformers' LinearAttentionLayer fields)."""
131
+ layer.conv_states[0], layer.recurrent_states[0] = conv, recurrent
132
+ layer.is_conv_states_initialized[0] = layer.is_recurrent_states_initialized[0] = True
133
+ layer.conv_kernel_size[0], layer.has_previous_state[0] = conv.shape[-1], previous
134
+ layer.device, layer.dtype = conv.device, conv.dtype
135
+
136
+
137
+ def is_attention(layer):
138
+ return isinstance(layer, DynamicLayer)
139
+
140
+
141
+ def layer_tensors(layer):
142
+ """The two tensors a prefix-cache layer holds: attention keys/values, or DeltaNet conv/recurrent states."""
143
+ return (layer.keys, layer.values) if is_attention(layer) else (layer.conv_states[0], layer.recurrent_states[0])
144
+
145
+
146
+ class Buffers:
147
+ """Per layer two flat buffers and the per-entry shape they hold: attention [heads, T, dim] (T set per view, `tokens`
148
+ entries x positions in total), DeltaNet states (`entries` entries)."""
149
+
150
+ def __init__(self, probe, tokens, entries, device):
151
+ self.slots = []
152
+ for layer in probe.layers:
153
+ ts, attention = layer_tensors(layer), is_attention(layer)
154
+ size = [tokens * t.shape[1] * t.shape[3] if attention else entries * t[0].numel() for t in ts]
155
+ self.slots.append((attention, [torch.zeros(n, dtype=t.dtype, device=device) for n, t in zip(size, ts)], [t.shape[1:] for t in ts]))
156
+
157
+ def views(self, n, length):
158
+ """Per layer, its two buffers viewed for n entries: attention [n, heads, length, dim], DeltaNet [n, *state]."""
159
+ def view(flat, shape, attention):
160
+ shape = (n, shape[0], length, shape[2]) if attention else (n, *shape)
161
+ return flat[:torch.Size(shape).numel()].view(shape)
162
+ return [tuple(view(f, s, attention) for f, s in zip(flats, shapes)) for attention, flats, shapes in self.slots]
163
+
164
+
165
+ class CudaGraphs:
166
+ def __init__(self, lm, pad_id):
167
+ self.lm, self.pad_id = lm, pad_id
168
+ self.device, self.dtype = next(lm.parameters()).device, next(lm.parameters()).dtype
169
+ self.pool, self.stream = torch.cuda.graph_pool_handle(), torch.cuda.Stream() # the graphs' shared memory pool, the capture stream
170
+ self.graphs = OrderedDict() # key -> (graph, input buffer)
171
+ self.pending = {} # key -> (pass body, input buffer): buckets that ran eagerly, not captured yet
172
+ self.eager_runs = {} # key -> how many passes of a pending bucket ran eagerly
173
+ self.failed = {} # key -> (why its capture failed, input buffer); the bucket keeps running eagerly
174
+ self.captures = 0
175
+ probe = DynamicCache(config=lm.config) # learn the cache layout (layer kinds, state shapes, dtypes) from one eager pass
176
+ with torch.no_grad():
177
+ lm(input_ids=torch.full((1, 16), pad_id, device=self.device), past_key_values=probe, use_cache=True)
178
+ # BufferKV and set_linear rely on this layout: plain attention layers and single-state DeltaNet layers that update
179
+ # their states in place. Anything else (another cache layer type, record_past, which assigns instead of copying)
180
+ # would leave the buffers stale and move probabilities silently, so refuse it here.
181
+ if not all(type(l) is DynamicLayer or (type(l) is LinearAttentionLayer and l.number_of_states == 1 and not l.record_past) for l in probe.layers):
182
+ raise ValueError(f"CUDA graphs support attention and single-state DeltaNet cache layers only; got {sorted({type(l).__name__ for l in probe.layers})}")
183
+ self.bank = Buffers(probe, GRAPH_STATES * BANK_WIDTH, GRAPH_STATES, self.device)
184
+ # the bank has one layout for every pass, so the state and row graphs agree on it: attention [GRAPH_STATES, heads,
185
+ # BANK_WIDTH, dim] with states right-aligned at the end, DeltaNet [GRAPH_STATES, *state]
186
+ self.bank_views = self.bank.views(GRAPH_STATES, BANK_WIDTH)
187
+ self.rowbuf = Buffers(probe, GRAPH_TOKENS, GRAPH_ROWS, self.device)
188
+ self.hidden = torch.zeros(GRAPH_ROWS * GRAPH_ROW * lm.config.hidden_size, dtype=self.dtype, device=self.device)
189
+
190
+ def _cache(self, views, filled, previous):
191
+ """A DynamicCache over buffer views: attention layers hold `filled` cached positions, DeltaNet layers the states."""
192
+ cache = DynamicCache(config=self.lm.config)
193
+ for i, (layer, v) in enumerate(zip(cache.layers, views)):
194
+ if is_attention(layer): cache.layers[i] = BufferKV(*v, filled)
195
+ else: set_linear(layer, *v, previous)
196
+ return cache
197
+
198
+ def _mask(self, allow):
199
+ """Additive attention mask [B, 1, Lq, Lk] from a boolean one [B, Lq, Lk], in the backbone's dtype."""
200
+ return torch.zeros(allow.shape, dtype=self.dtype, device=self.device).masked_fill(~allow, torch.finfo(self.dtype).min)[:, None]
201
+
202
+ def _forward(self, ids, pos, full_mask, linear_mask, cache):
203
+ return self.lm(input_ids=ids, position_ids=pos, attention_mask={"full_attention": full_mask, "linear_attention": linear_mask},
204
+ past_key_values=cache, use_cache=True).last_hidden_state
205
+
206
+ def _replay(self, key, body, rows):
207
+ """One pass for `key` on `rows` (token ids, positions, lengths, indices; int64): the graph's replay, or body() run
208
+ eagerly when the bucket has no graph yet (it then waits in `pending`). body(buf) runs the pass reading the uploaded
209
+ rows from buf. Every body for a key reads and writes the same buffer views, so the one kept in `pending` stands for
210
+ all of them."""
211
+ if key in self.graphs:
212
+ self.graphs.move_to_end(key)
213
+ graph, buf = self.graphs[key]
214
+ elif key in self.failed:
215
+ graph, buf = None, self.failed[key][1]
216
+ else:
217
+ graph, buf = None, self.pending.setdefault(key, (body, torch.zeros((len(rows), len(rows[0])), dtype=torch.long, device=self.device)))[1]
218
+ self.eager_runs[key] = self.eager_runs.get(key, 0) + 1
219
+ buf.copy_(torch.tensor(rows, dtype=torch.long), non_blocking=True)
220
+ if graph is not None: graph.replay()
221
+ else: body(buf)
222
+
223
+ @torch.no_grad()
224
+ def capture_pending(self, limit=None):
225
+ """Capture graphs for up to `limit` (None = all) buckets that have run eagerly, the most-run first. The caller must
226
+ keep other passes out (kev.serve captures on its model thread). Each capture overwrites the shared buffers, which is
227
+ fine between passes: a pass refills them. On a side stream, a warm-up pass first (Triton autotuning and cuBLAS
228
+ setup must not happen inside a capture), then the capture; not torch.cuda.graph, whose synchronize, gc.collect and
229
+ empty_cache on entry would stall a busy server for each capture. thread_local: CUDA calls of other threads (a
230
+ request tokenizing, a model card read) cannot invalidate the capture. A capture that fails leaves its bucket running
231
+ eagerly, for good."""
232
+ for _ in range(len(self.pending) if limit is None else min(limit, len(self.pending))):
233
+ key = max(self.pending, key=self.eager_runs.get)
234
+ (body, buf), _ = self.pending.pop(key), self.eager_runs.pop(key)
235
+ current, graph = torch.cuda.current_stream(), torch.cuda.CUDAGraph()
236
+ self.stream.wait_stream(current)
237
+ try:
238
+ with torch.cuda.stream(self.stream):
239
+ body(buf) # warm-up: autotuning and lazy setup happen outside the capture
240
+ graph.capture_begin(pool=self.pool, capture_error_mode="thread_local")
241
+ try: body(buf)
242
+ finally: graph.capture_end()
243
+ except Exception as e: # e.g. out of memory for a new shape: serve it eagerly rather than stop serving
244
+ # a failed capture_end leaves the allocator routing this thread's allocations into the graph pool, which
245
+ # would fail every later capture and let eager passes allocate graph memory: end that explicitly
246
+ try: getattr(torch._C, "_cuda_endAllocateToPool", lambda *a: None)(self.device.index, self.pool) # private API: best effort
247
+ except Exception: pass
248
+ self.failed[key] = (f"{type(e).__name__}: {e}", buf)
249
+ print(f"kev.cuda_graphs: capturing {key} failed, it runs eagerly: {self.failed[key][0]}", flush=True)
250
+ continue
251
+ finally:
252
+ current.wait_stream(self.stream)
253
+ self.graphs[key] = (graph, buf); self.captures += 1
254
+ while len(self.graphs) > GRAPHS_KEPT: self.graphs.popitem(last=False)
255
+
256
+ def capture_due(self, idle):
257
+ """Capture one pending graph now? Always when idle; under load once a bucket keeps running eagerly (HOT_BUCKET
258
+ eager passes), because a request arriving meanwhile waits for the capture (~0.4 s)."""
259
+ runs = list(self.eager_runs.values()) # a snapshot: wait_idle and /v1/models call this off the model thread
260
+ return bool(runs) and (idle or max(runs) >= HOT_BUCKET)
261
+
262
+ def stats(self):
263
+ return {"captured": self.captures, "kept": len(self.graphs), "pending": len(self.pending), "failed": len(self.failed)}
264
+
265
+ # Serving: which requests the graphed passes take, and one batch of them
266
+
267
+ @staticmethod
268
+ def admits(state_len, row_lens, cached):
269
+ """Whether a request's passes fit the graphed passes: its rows, and its state (a new one runs the graphed state
270
+ pass; a cached one only has to fit the bank)."""
271
+ return bucket(max(row_lens)) <= GRAPH_ROW and (state_len <= BANK_WIDTH if cached else bucket(state_len) <= GRAPH_STATE)
272
+
273
+ @torch.no_grad()
274
+ def run(self, requests):
275
+ """Admitted requests (Request) -> (their picked hidden states, float32 [sum of picks, d] in request, row, pick
276
+ order; per request its state's cache: the cached one, a new one if `keep`, else None). States grouped by length
277
+ (length_groups) share a state pass (identical states once); each request takes one bank entry, so at most GRAPH_STATES per group.
278
+ Every row pass writes its picks straight to their place in the output."""
279
+ at = [0]
280
+ for r in requests: at.append(at[-1] + sum(map(len, r.picks)))
281
+ out = torch.empty((at[-1], self.lm.config.hidden_size), dtype=torch.float32, device=self.device)
282
+ caches = [None] * len(requests)
283
+ for group in length_groups([len(r.state) for r in requests], GRAPH_STATES):
284
+ new = {} # distinct new states of the group -> (ids, positions, keep)
285
+ for i in group:
286
+ r = requests[i]
287
+ if r.cached is None: new[tuple(r.state)] = (r.state, r.positions, r.keep or new.get(tuple(r.state), (0, 0, False))[2])
288
+ made = dict(zip(new, self.states(list(new.values()), bucket(max(map(len, new)))))) if new else {}
289
+ entry = {S: j for j, S in enumerate(new)} # cached states go into the entries after the new ones
290
+ rows, free = [], len(new)
291
+ for i in group: # after the state pass: its padded entries are written
292
+ r = requests[i]
293
+ if r.cached is None: e, caches[i] = entry[tuple(r.state)], made[tuple(r.state)]
294
+ else: e, free, caches[i] = free, free + 1, r.cached; self.load_state(e, r.cached, len(r.state))
295
+ offset = at[i]
296
+ for (ids, pos), picks in zip(r.rows, r.picks):
297
+ rows.append(_Row(ids, pos, e, len(r.state), picks, offset)); offset += len(picks)
298
+ self.rows(rows, out)
299
+ return out, caches
300
+
301
+ @torch.no_grad()
302
+ def states(self, items, Sb):
303
+ """State pass for up to GRAPH_STATES states, items = [(ids, pos, keep)] each at most Sb <= GRAPH_STATE tokens.
304
+ Leaves state i in bank entry i and returns, for each kept state, a DynamicCache equal to the eager prefix pass's
305
+ (None for the others: copying every state of a batch cost more than the passes)."""
306
+ Nb = count_bucket(len(items))
307
+
308
+ def body(buf): # buf rows = [ids | positions | real length], left-padded
309
+ pad = Sb - buf[:, 2 * Sb:]
310
+ i = torch.arange(Sb, device=self.device)
311
+ q, k = i[None, :, None], i[None, None, :]
312
+ allow = ((k <= q) & (k >= pad[:, :, None])) | (k == q)
313
+ views = [(a[:Nb, :, BANK_WIDTH - Sb:], b[:Nb, :, BANK_WIDTH - Sb:]) if attention else (a[:Nb], b[:Nb])
314
+ for (attention, *_), (a, b) in zip(self.bank.slots, self.bank_views)]
315
+ self._forward(buf[:, :Sb], buf[:, Sb:2 * Sb], self._mask(allow), (i[None] >= pad).long(), self._cache(views, 0, False))
316
+
317
+ rows = [[self.pad_id] * (Sb - len(ids)) + list(ids) + [0] * (Sb - len(ids)) + list(pos) + [len(ids)] for ids, pos, _ in items]
318
+ self._replay(("states", Nb, Sb), body, rows + [[self.pad_id] * Sb + [0] * Sb + [0]] * (Nb - len(items)))
319
+ out = [] # a cache for each kept state: its own copy of the bank entry (the bank belongs to the next pass)
320
+ for j, (ids, _, keep) in enumerate(items):
321
+ if not keep: out.append(None); continue
322
+ cache, S = DynamicCache(config=self.lm.config), len(ids)
323
+ for layer, (a, b) in zip(cache.layers, self.bank_views):
324
+ if is_attention(layer):
325
+ layer.keys, layer.values = a[j:j + 1, :, BANK_WIDTH - S:].clone(), b[j:j + 1, :, BANK_WIDTH - S:].clone()
326
+ layer.dtype, layer.device, layer.is_initialized = self.dtype, self.device, True
327
+ else:
328
+ set_linear(layer, a[j:j + 1].clone(), b[j:j + 1].clone(), True)
329
+ out.append(cache)
330
+ return out
331
+
332
+ @torch.no_grad()
333
+ def load_state(self, entry, cache, S):
334
+ """Copy a cached state (a DynamicCache of S <= BANK_WIDTH tokens, from either path) into bank entry `entry`."""
335
+ for layer, v in zip(cache.layers, self.bank_views):
336
+ for buf, t in zip(v, layer_tensors(layer)):
337
+ if is_attention(layer): buf[entry, :, BANK_WIDTH - S:].copy_(t[0, :, t.shape[-2] - S:])
338
+ else: buf[entry].copy_(t[0])
339
+
340
+ def rows(self, rows, out):
341
+ """Row passes for question rows (_Row), each continuing its state in the bank; each row's picked hidden states land
342
+ in out[row.at:]. The rows see the last Sr positions of the bank, Sr the longest state rounded up to a power of two
343
+ (it only lengthens attention). Rows are grouped by length (length_groups), then split to fit the buffers."""
344
+ for idx in length_groups([len(r.ids) for r in rows], GRAPH_ROWS):
345
+ Lb, Sr = bucket(max(len(rows[i].ids) for i in idx)), pow2(max(16, max(rows[i].state_len for i in idx)))
346
+ group = min(GRAPH_ROWS, GRAPH_TOKENS // (Sr + Lb))
347
+ group = max(n for n in range(1, group + 1) if count_bucket(n) <= group) # the most rows whose padded count fits
348
+ for start in range(0, len(idx), group):
349
+ part = [rows[i] for i in idx[start:start + group]]
350
+ h = self._row_pass(part, Sr, Lb)
351
+ src, dst = zip(*[(n * Lb + p, r.at + j) for n, r in enumerate(part) for j, p in enumerate(r.picks)])
352
+ src, dst = torch.tensor([src, dst]).to(self.device, non_blocking=True)
353
+ out.index_copy_(0, dst, h.flatten(0, 1).index_select(0, src).float()) # only the wanted positions leave the pass buffer
354
+
355
+ def _row_pass(self, rows, Sr, Lb):
356
+ """One row pass; -> the pass buffer [rows, Lb, d] (valid until the next pass)."""
357
+ Nb, T = count_bucket(len(rows)), Sr + Lb
358
+ hidden = self.hidden[:Nb * Lb * self.lm.config.hidden_size].view(Nb, Lb, -1)
359
+
360
+ def body(buf): # buf rows = [ids | positions | row length | state length | bank entry]
361
+ rowlen, plen, entry = buf[:, 2 * Lb:2 * Lb + 1], buf[:, 2 * Lb + 1:2 * Lb + 2], buf[:, 2 * Lb + 2]
362
+ views = self.rowbuf.views(Nb, T)
363
+ for (attention, *_), dst, src in zip(self.bank.slots, views, self.bank_views): # each row's state, from the bank
364
+ for d, s in zip(dst, src):
365
+ if attention: d[:, :, :Sr].copy_(s[:, :, BANK_WIDTH - Sr:].index_select(0, entry))
366
+ else: d.copy_(s.index_select(0, entry))
367
+ q = torch.arange(Lb, device=self.device)[None, :, None]
368
+ k = torch.arange(T, device=self.device)[None, None, :]
369
+ allow = ((k < Sr) & (k >= Sr - plen[:, :, None])) | ((k >= Sr) & (k - Sr <= q) & (k - Sr < rowlen[:, :, None])) | (k - Sr == q)
370
+ linear = (torch.arange(Lb, device=self.device)[None] < rowlen).long()
371
+ hidden.copy_(self._forward(buf[:, :Lb], buf[:, Lb:2 * Lb], self._mask(allow), linear, self._cache(views, Sr, True)))
372
+
373
+ uploads = [list(r.ids) + [self.pad_id] * (Lb - len(r.ids)) + list(r.pos) + [0] * (Lb - len(r.pos)) + [len(r.ids), r.state_len, r.entry] for r in rows]
374
+ self._replay(("rows", Nb, Lb, Sr), body, uploads + [[self.pad_id] * Lb + [0] * Lb + [0, 0, 0]] * (Nb - len(rows)))
375
+ return hidden[:len(rows)]
runtime/kev/data.py ADDED
@@ -0,0 +1,410 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Convert public labelled datasets into TypeSafe-shaped requests, then through api.to_record() so the
2
+ training format is byte-identical to what /v1/systemone feeds the model.
3
+
4
+ A "labelled request" is {"state": JSONContent, "questions": {id: {type, instructions, criteria, "label": ..., "src": str}}}
5
+ label: choice -> option key, noul -> bool, score -> level index.
6
+ materialize() -> internal record {"state": str, "questions": [{"instr", "options", "label": int, "src"}]}
7
+ """
8
+ import hashlib
9
+ import json
10
+ from pathlib import Path
11
+ import random
12
+ from datasets import load_dataset
13
+ from .api import SystemOneRequest, to_record
14
+
15
+
16
+ REPOS = {"banking77": "legacy-datasets/banking77", "boolq": "google/boolq", "agnews": "fancyzhx/ag_news",
17
+ "mnli": "nyu-mll/multi_nli", "sst5": "SetFit/sst5", "yelp": "Yelp/yelp_review_full"}
18
+
19
+
20
+ def source_seed(seed, source):
21
+ return int.from_bytes(hashlib.sha256(f"{seed}:{source}".encode()).digest()[:8], "big")
22
+
23
+
24
+ # Repos whose main branch is a (no longer supported) loading script; read the Hub's auto-converted parquet branch instead.
25
+ PARQUET_BRANCH = {"CogComp/trec": "refs/convert/parquet"}
26
+
27
+
28
+ def dataset_ref(repo):
29
+ """(repo, revision-to-pin) for a repo id; parquet-branch repos pin that branch's sha, not main."""
30
+ return repo.partition(":")[0], PARQUET_BRANCH.get(repo.partition(":")[0])
31
+
32
+
33
+ class Source(random.Random):
34
+ """One source's sampling context: a seeded RNG (the converters draw everything from it), the dataset revision to pin,
35
+ and the provenance of every row drawn (`origins`, filled by _sample so build() can attach it to each record)."""
36
+
37
+ def __init__(self, seed, revision=None):
38
+ super().__init__(seed)
39
+ self.revision = revision
40
+ self.origins = []
41
+
42
+
43
+ def _dataset(repo, split, src):
44
+ """repo may be 'owner/name' or 'owner/name:config'; the Source carries the revision to pin."""
45
+ name, _, config = repo.partition(":")
46
+ return load_dataset(name, config or None, split=split, revision=src.revision or PARQUET_BRANCH.get(name))
47
+
48
+ NONE = "None of the above"
49
+ # "None of the above" options must appear both as the correct answer and as a wrong alternative, with varied
50
+ # wording, or the model learns "this wording => pick it" (it did, in the first training run).
51
+ NONE_OPTIONS = [("other", "None of the above"), ("other", "A reason that fits none of the above"), ("none", "None of these"),
52
+ ("other", "Something else"), ("not_listed", "Not listed here"), ("none_of_the_above", None),
53
+ ("other", "A category that fits none of the above"), ("other", "None of the listed options apply"),
54
+ ("unknown", "Cannot be determined from the options given"), ("other", "Other"), ("none", None),
55
+ ("other", "An answer not covered by the other options"), ("no_match", "No option matches")]
56
+ DISTRACTORS = {"weather": "Bad weather caused it", "purple": "The colour purple", "pancakes": "A recipe for pancakes", "taxes": "Unrelated: quarterly tax filing"}
57
+
58
+ AG = {"world": "World news: politics, international affairs, conflicts", "sports": "Sports: games, athletes, teams, results",
59
+ "business": "Business: companies, markets, economy, finance", "scitech": "Science and technology: research, gadgets, software, space"}
60
+ MNLI = {"entailment": "The hypothesis follows from the premise", "neutral": "The hypothesis may or may not be true given the premise", "contradiction": "The hypothesis contradicts the premise"}
61
+ SST5 = ["very negative", "negative", "neutral", "positive", "very positive"]
62
+ YELP = ["1 star: terrible experience", "2 stars: poor", "3 stars: average", "4 stars: good", "5 stars: excellent"]
63
+ BANK_TEMPLATES = ["Customer asks about {}", "Issue concerning {}", "Request related to {}", "{}"]
64
+
65
+
66
+ def _wrap_state(text, rng):
67
+ r = rng.random()
68
+ if r < 0.15: return {"document": text}
69
+ if r < 0.25: return {"ticket": {"channel": rng.choice(["email", "chat", "web form"]), "body": text}}
70
+ if r < 0.32: return [{"role": "customer", "content": text}]
71
+ return text
72
+
73
+
74
+ def _instr(text, rng):
75
+ return {"question": text, "focus": rng.choice(["Use only the information given.", "Pick the single best fit.", "Consider the whole message."])} if rng.random() < 0.15 else text
76
+
77
+
78
+ def _desc(desc, rng, p_null=0.3, p_struct=0.1):
79
+ r = rng.random()
80
+ if r < p_null: return None
81
+ if r < p_null + p_struct: return {"what": desc}
82
+ return desc
83
+
84
+
85
+ def _sample(ds, n, src):
86
+ rows = []
87
+ src.origins = []
88
+ for i in src.sample(range(len(ds)), min(n, len(ds))):
89
+ row = ds[i]
90
+ if row.get("label", 0) == -1:
91
+ continue
92
+ text = next((row[k] for k in ("text", "premise", "passage", "content", "question", "sentence") if isinstance(row.get(k), str)), json.dumps(row, sort_keys=True))
93
+ normalized = " ".join(text.casefold().split())
94
+ src.origins.append({"row": i, "text_sha256": hashlib.sha256(normalized.encode()).hexdigest(),
95
+ "row_sha256": hashlib.sha256(json.dumps(row, sort_keys=True).encode()).hexdigest()})
96
+ rows.append(row)
97
+ return rows
98
+
99
+
100
+ def _banking(split, n, rng):
101
+ ds = _dataset("legacy-datasets/banking77", split=split, src=rng)
102
+ names = ds.features["label"].names
103
+ out = []
104
+ for ex in _sample(ds, n, rng):
105
+ t = rng.choice(BANK_TEMPLATES)
106
+ crit = {k: _desc(t.format(k.replace("_", " ")), rng, p_null=0.5, p_struct=0.0) for k in names}
107
+ out.append({"state": _wrap_state(ex["text"], rng), "questions": {"intent": {"type": "choice", "instructions": _instr("Which banking intent best describes this customer message?", rng), "criteria": crit, "label": names[ex["label"]], "src": "banking77"}}})
108
+ return out
109
+
110
+
111
+ def _boolq(split, n, rng):
112
+ ds = _dataset("google/boolq", split=split, src=rng)
113
+ out = []
114
+ for ex in _sample(ds, n, rng):
115
+ q = {"type": "noul", "instructions": _instr(ex["question"].strip().rstrip("?") + "?", rng), "label": bool(ex["answer"]), "src": "boolq"}
116
+ if rng.random() < 0.4: q["criteria"] = {"true": "The passage supports a yes answer", "false": "The passage supports a no answer or does not say"}
117
+ out.append({"state": _wrap_state(ex["passage"], rng), "questions": {"answer": q}})
118
+ return out
119
+
120
+
121
+ def _agnews(split, n, rng):
122
+ ds = _dataset("fancyzhx/ag_news", split=split, src=rng)
123
+ keys = list(AG)
124
+ out = []
125
+ for ex in _sample(ds, n, rng):
126
+ y = keys[ex["label"]]
127
+ qs = {"topic": {"type": "choice", "instructions": _instr("What is the topic of this article?", rng), "criteria": {k: _desc(v, rng) for k, v in AG.items()}, "label": y, "src": "agnews"}}
128
+ for k in rng.sample(keys, 2):
129
+ qs[f"is_{k}"] = {"type": "noul", "instructions": f"Is this article about {AG[k].split(':')[0].lower()}?", "label": k == y, "src": "agnews_yn"}
130
+ out.append({"state": _wrap_state(ex["text"], rng), "questions": qs})
131
+ return out
132
+
133
+
134
+ def _mnli(split, n, rng):
135
+ ds = _dataset("nyu-mll/multi_nli", split=split, src=rng)
136
+ keys = list(MNLI)
137
+ return [{"state": _wrap_state(ex["premise"], rng), "questions": {"relation": {"type": "choice", "instructions": _instr(f'Hypothesis: "{ex["hypothesis"]}" How does it relate to the premise?', rng), "criteria": {k: _desc(v, rng) for k, v in MNLI.items()}, "label": keys[ex["label"]], "src": "mnli"}}}
138
+ for ex in _sample(ds, n, rng) if ex["label"] >= 0]
139
+
140
+
141
+ def _sst5(split, n, rng):
142
+ ds = _dataset("SetFit/sst5", split=split, src=rng)
143
+ return [{"state": _wrap_state(ex["text"], rng), "questions": {"sentiment": {"type": "score", "instructions": _instr("What is the sentiment of this review sentence?", rng), "criteria": list(SST5), "label": ex["label"], "src": "sst5"}}} for ex in _sample(ds, n, rng)]
144
+
145
+
146
+ def _yelp(split, n, rng):
147
+ ds = _dataset("Yelp/yelp_review_full", split=split, src=rng)
148
+ out = []
149
+ for ex in _sample(ds, n, rng):
150
+ text = " ".join(ex["text"].split()[:220])
151
+ qs = {"rating": {"type": "score", "instructions": _instr("How many stars did this reviewer give?", rng), "criteria": list(YELP), "label": ex["label"], "src": "yelp"},
152
+ "recommend": {"type": "noul", "instructions": "Would this reviewer recommend the business?", "criteria": {"true": "Clearly positive overall", "false": "Negative or mixed"}, "label": ex["label"] >= 3, "src": "yelp_yn"}}
153
+ out.append({"state": _wrap_state(text, rng), "questions": qs})
154
+ return out
155
+
156
+
157
+ SOURCES = {"banking77": (_banking, "train", "test"), "boolq": (_boolq, "train", "validation"), "agnews": (_agnews, "train", "test"),
158
+ "mnli": (_mnli, "train", "validation_matched"), "sst5": (_sst5, "train", "test"), "yelp": (_yelp, "train", "test")}
159
+
160
+
161
+ # Eval-only sources for transfer measurement. Never passed to training; kept out of SOURCES so build() defaults cannot
162
+ # pick them up. Rotten Tomatoes (SST parent) and SNLI (MNLI sibling) are deliberately excluded.
163
+ TREC = {"abbreviation": "Asks what an abbreviation stands for", "entity": "Asks about a thing, object, animal, product, or creative work",
164
+ "description": "Asks for a definition, description, reason, or manner", "human": "Asks about a person, group, or organisation",
165
+ "location": "Asks about a place", "number": "Asks for a number, date, count, or other numeric value"}
166
+ EMOTION = {"sadness": None, "joy": None, "love": None, "anger": None, "fear": None, "surprise": None}
167
+ AMAZON = ["1 star: very negative", "2 stars: negative", "3 stars: mixed", "4 stars: positive", "5 stars: very positive"]
168
+
169
+
170
+ def _trec(split, n, rng):
171
+ ds = _dataset("CogComp/trec", split=split, src=rng)
172
+ keys = list(TREC)
173
+ return [{"state": _wrap_state(ex["text"], rng), "questions": {"answer_type": {"type": "choice", "instructions": "What kind of answer does this question ask for?",
174
+ "criteria": dict(TREC), "label": keys[ex["coarse_label"]], "src": "trec"}}} for ex in _sample(ds, n, rng)]
175
+
176
+
177
+ def _dbpedia(split, n, rng):
178
+ ds = _dataset("fancyzhx/dbpedia_14", split=split, src=rng)
179
+ names = [x.lower().replace(" ", "_") for x in ds.features["label"].names]
180
+ return [{"state": _wrap_state(" ".join(ex["content"].split()[:200]), rng), "questions": {"category": {"type": "choice", "instructions": "Which category does the subject of this encyclopedia text belong to?",
181
+ "criteria": {k: None for k in names}, "label": names[ex["label"]], "src": "dbpedia14"}}} for ex in _sample(ds, n, rng)]
182
+
183
+
184
+ def _emotion(split, n, rng):
185
+ ds = _dataset("dair-ai/emotion:split", split=split, src=rng)
186
+ keys = list(EMOTION)
187
+ return [{"state": ex["text"], "questions": {"emotion": {"type": "choice", "instructions": "Which emotion does the writer express?",
188
+ "criteria": dict(EMOTION), "label": keys[ex["label"]], "src": "emotion"}}} for ex in _sample(ds, n, rng)]
189
+
190
+
191
+ def _imdb(split, n, rng):
192
+ ds = _dataset("stanfordnlp/imdb", split=split, src=rng)
193
+ return [{"state": _wrap_state(" ".join(ex["text"].replace("<br />", " ").split()[:220]), rng), "questions": {"positive": {"type": "noul", "instructions": "Is this movie review positive?",
194
+ "criteria": {"true": "The reviewer liked the film overall", "false": "The reviewer disliked the film overall"}, "label": ex["label"] == 1, "src": "imdb"}}} for ex in _sample(ds, n, rng)]
195
+
196
+
197
+ def _amazon(split, n, rng):
198
+ ds = _dataset("SetFit/amazon_reviews_multi_en", split=split, src=rng)
199
+ return [{"state": _wrap_state(" ".join(ex["text"].split()[:220]), rng), "questions": {"stars": {"type": "score", "instructions": "How many stars did this product reviewer give?",
200
+ "criteria": list(AMAZON), "label": ex["label"], "src": "amazon"}}} for ex in _sample(ds, n, rng)]
201
+
202
+
203
+ def _qnli(split, n, rng):
204
+ ds = _dataset("nyu-mll/glue:qnli", split=split, src=rng)
205
+ return [{"state": _wrap_state(ex["sentence"], rng), "questions": {"answers": {"type": "noul", "instructions": f'Does the sentence contain the answer to this question: "{ex["question"]}"',
206
+ "label": ex["label"] == 0, "src": "qnli"}}} for ex in _sample(ds, n, rng)]
207
+
208
+
209
+ def _offensive(split, n, rng):
210
+ ds = _dataset("cardiffnlp/tweet_eval:offensive", split=split, src=rng)
211
+ return [{"state": ex["text"], "questions": {"offensive": {"type": "noul", "instructions": "Is this post offensive?",
212
+ "criteria": {"true": "Contains insults, threats, profanity directed at someone, or hateful content", "false": "Not offensive"},
213
+ "label": ex["label"] == 1, "src": "tweet_offensive"}}} for ex in _sample(ds, n, rng)]
214
+
215
+
216
+ def _mmlu(split, n, rng):
217
+ ds = _dataset("cais/mmlu:all", split=split, src=rng)
218
+ out = []
219
+ for ex in _sample(ds, n, rng):
220
+ keys = ["a", "b", "c", "d"]
221
+ out.append({"state": {"subject": ex["subject"].replace("_", " "), "question": ex["question"]},
222
+ "questions": {"answer": {"type": "choice", "instructions": "Which option correctly answers the question?",
223
+ "criteria": dict(zip(keys, ex["choices"])), "label": keys[ex["answer"]], "src": "mmlu"}}})
224
+ return out
225
+
226
+
227
+ def _paws(split, n, rng):
228
+ ds = _dataset("google-research-datasets/paws:labeled_final", split=split, src=rng)
229
+ return [{"state": _wrap_state(ex["sentence1"], rng), "questions": {"paraphrase": {"type": "noul", "instructions": f'Does this sentence mean the same thing: "{ex["sentence2"]}"',
230
+ "criteria": {"true": "Same meaning, possibly reworded", "false": "Different meaning, even if most words match"}, "label": ex["label"] == 1, "src": "paws"}}}
231
+ for ex in _sample(ds, n, rng)]
232
+
233
+
234
+ def _sciq(split, n, rng):
235
+ ds = _dataset("allenai/sciq", split=split, src=rng)
236
+ out = []
237
+ for ex in _sample(ds, n, rng):
238
+ options = [ex["correct_answer"], ex["distractor1"], ex["distractor2"], ex["distractor3"]]
239
+ keys = ["a", "b", "c", "d"]; order = list(range(4)); rng.shuffle(order)
240
+ crit = {keys[i]: options[j] for i, j in enumerate(order)}
241
+ out.append({"state": {"passage": ex["support"], "question": ex["question"]} if ex["support"] else {"question": ex["question"]},
242
+ "questions": {"answer": {"type": "choice", "instructions": "Which option answers the science question?", "criteria": crit,
243
+ "label": keys[order.index(0)], "src": "sciq"}}})
244
+ return out
245
+
246
+
247
+ def _mcq(ex_q, labels, texts, answer_label, src, rng, state_extra=None):
248
+ """Knowledge MCQ -> Choice with neutral keys; option order shuffled per record so keys carry no information."""
249
+ order = list(range(len(texts))); rng.shuffle(order)
250
+ keys = [f"opt_{i + 1}" for i in range(len(texts))]
251
+ crit = {keys[i]: texts[j] for i, j in enumerate(order)}
252
+ label = keys[order.index(labels.index(answer_label))]
253
+ state = {"question": ex_q}
254
+ if state_extra: state.update(state_extra)
255
+ return {"state": state, "questions": {"answer": {"type": "choice", "instructions": "Which option correctly answers the question?", "criteria": crit, "label": label, "src": src}}}
256
+
257
+
258
+ def _arc(split, n, rng):
259
+ ds = _dataset("allenai/ai2_arc:ARC-Challenge", split=split, src=rng)
260
+ return [_mcq(ex["question"], list(ex["choices"]["label"]), list(ex["choices"]["text"]), ex["answerKey"], "arc", rng) for ex in _sample(ds, n, rng)]
261
+
262
+
263
+ def _openbookqa(split, n, rng):
264
+ ds = _dataset("allenai/openbookqa:main", split=split, src=rng)
265
+ return [_mcq(ex["question_stem"], list(ex["choices"]["label"]), list(ex["choices"]["text"]), ex["answerKey"], "openbookqa", rng) for ex in _sample(ds, n, rng)]
266
+
267
+
268
+ def _csqa(split, n, rng):
269
+ ds = _dataset("tau/commonsense_qa", split=split, src=rng)
270
+ return [_mcq(ex["question"], list(ex["choices"]["label"]), list(ex["choices"]["text"]), ex["answerKey"], "csqa", rng) for ex in _sample(ds, n, rng) if ex["answerKey"]]
271
+
272
+
273
+ ALL_SOURCES = {**SOURCES, "trec": (_trec, "train", "test"), "dbpedia14": (_dbpedia, "train", "test"), "emotion": (_emotion, "train", "test"),
274
+ "imdb": (_imdb, "train", "test"), "amazon": (_amazon, "train", "test"), "qnli": (_qnli, "train", "validation"),
275
+ "tweet_offensive": (_offensive, "train", "test"), "mmlu": (_mmlu, "test", "test"),
276
+ "paws": (_paws, "train", "test"), "sciq": (_sciq, "train", "test"),
277
+ "arc": (_arc, "train", "test"), "openbookqa": (_openbookqa, "train", "test"), "csqa": (_csqa, "train", "validation")}
278
+ ALL_REPOS = {**REPOS, "trec": "CogComp/trec", "dbpedia14": "fancyzhx/dbpedia_14", "emotion": "dair-ai/emotion", "imdb": "stanfordnlp/imdb",
279
+ "amazon": "SetFit/amazon_reviews_multi_en", "qnli": "nyu-mll/glue", "tweet_offensive": "cardiffnlp/tweet_eval", "mmlu": "cais/mmlu",
280
+ "paws": "google-research-datasets/paws", "sciq": "allenai/sciq",
281
+ "arc": "allenai/ai2_arc", "openbookqa": "allenai/openbookqa", "csqa": "tau/commonsense_qa"}
282
+
283
+ # Policy (PLAN.md step 1). A source is trainable or eval-only; suites record both lists and training refuses eval-only
284
+ # sources. MMLU is a knowledge probe and stays eval-only permanently; Emotion/TweetEval are noisy-label honesty checks;
285
+ # QNLI/PAWS/SciQ measure reading transfer. Rotten Tomatoes (SST parent) and SNLI (MNLI sibling) are excluded entirely.
286
+ # Knowledge MCQ (ARC-Challenge, OpenBookQA, CommonsenseQA) is trainable: the hypothesis (PLAN.md, overnight) is that the
287
+ # pointer readout under-uses the base model's knowledge (8B scores 0.65 on 4-way MMLU, below its base-model level) and
288
+ # that a small MCQ mix teaches the head to tap it. MMLU and SciQ stay eval-only; ARC/SciQ are distinct datasets.
289
+ TRAINABLE = ("banking77", "boolq", "agnews", "mnli", "sst5", "yelp", "trec", "dbpedia14", "amazon", "imdb", "arc", "openbookqa", "csqa")
290
+ EVAL_ONLY = ("mmlu", "emotion", "tweet_offensive", "qnli", "paws", "sciq")
291
+ assert set(TRAINABLE) | set(EVAL_ONLY) == set(ALL_SOURCES) and not set(TRAINABLE) & set(EVAL_ONLY)
292
+
293
+ # transfer-v1 (frozen before the policy existed) used these eight; kept so its manifest can be re-derived.
294
+ TRANSFER_SOURCES = {k: ALL_SOURCES[k] for k in ("trec", "dbpedia14", "emotion", "imdb", "amazon", "qnli", "tweet_offensive", "mmlu")}
295
+ TRANSFER_REPOS = {k: ALL_REPOS[k] for k in TRANSFER_SOURCES}
296
+
297
+
298
+ def build(n_per_source, split="train", seed=0, exclude=(), only=(), revisions=None, sources=None, repos=None):
299
+ sources = SOURCES if sources is None else sources
300
+ repos = REPOS if repos is None else repos
301
+ unknown = (set(exclude) | set(only)) - sources.keys()
302
+ if unknown:
303
+ raise ValueError(f"unknown sources: {sorted(unknown)}")
304
+ reqs = []
305
+ for name, (fn, tr, te) in sources.items():
306
+ if name in exclude or (only and name not in only): continue
307
+ src = Source(source_seed(seed, name), (revisions or {}).get(repos[name]))
308
+ source_split = tr if split == "train" else te
309
+ records = fn(source_split, n_per_source, src)
310
+ if len(records) != len(src.origins):
311
+ raise ValueError(f"provenance mismatch for {name}")
312
+ for record, origin in zip(records, src.origins):
313
+ record["_meta"] = {**origin, "source": name, "repo": repos[name], "revision": src.revision,
314
+ "split": source_split, "id": f"{name}/{source_split}/{origin['row']}"}
315
+ reqs.extend(records)
316
+ random.Random(seed).shuffle(reqs)
317
+ return reqs
318
+
319
+
320
+ def augment(req, rng, p_none=0.1, p_none_distract=0.12, p_distract=0.15):
321
+ """Choice only: permute option order (always); sometimes add a 'none of the above' option, either as the correct
322
+ answer (true option removed) or as a wrong alternative (true option kept); sometimes add an irrelevant distractor."""
323
+ if min(p_none, p_none_distract, p_distract) < 0 or p_none + p_none_distract + p_distract > 1:
324
+ raise ValueError("augmentation probabilities must be nonnegative and sum to at most one")
325
+ out = {"state": req["state"], "questions": {}}
326
+ for qid, q in req["questions"].items():
327
+ if q["type"] != "choice":
328
+ out["questions"][qid] = q; continue
329
+ crit, y = dict(q["criteria"]), q["label"]
330
+ if q.get("target") is not None: # soft-target questions: permute only; inserting or swapping options would change the target's meaning
331
+ keys = list(crit); rng.shuffle(keys); out["questions"][qid] = {**q, "criteria": {k: crit[k] for k in keys}}; continue
332
+ r = rng.random()
333
+ none_options = [(k, v) for k, v in NONE_OPTIONS if k not in crit]
334
+ distractors = [k for k in DISTRACTORS if k not in crit]
335
+ if len(crit) > 2 and r < p_none and none_options:
336
+ nk, nd = rng.choice(none_options); crit.pop(y); crit[nk] = nd; y = nk
337
+ elif p_none <= r < p_none + p_none_distract and len(crit) < 255 and none_options:
338
+ nk, nd = rng.choice(none_options); crit[nk] = nd
339
+ elif p_none + p_none_distract <= r < p_none + p_none_distract + p_distract and len(crit) < 255 and distractors:
340
+ k = rng.choice(distractors); crit[k] = DISTRACTORS[k]
341
+ keys = list(crit); rng.shuffle(keys)
342
+ out["questions"][qid] = {**q, "criteria": {k: crit[k] for k in keys}, "label": y}
343
+ return out
344
+
345
+
346
+ def none_pair(req, rng):
347
+ """Minimal pair for the none-of-the-above shortcut: the same state and question rendered twice, once with the true
348
+ option present (a none option is wrong) and once with it removed (the same none option is right). Everything else,
349
+ including option order, is identical, so the only difference the model can use is whether the evidence matches an
350
+ option. Returns [] when the request has no eligible Choice (>=3 options, no none key already present, no soft target:
351
+ removing the labelled option or adding a none option would change what the target means, as in augment)."""
352
+ eligible = [(qid, q) for qid, q in req["questions"].items() if q["type"] == "choice" and len(q["criteria"]) >= 3 and q.get("target") is None]
353
+ if not eligible: return []
354
+ qid, q = rng.choice(eligible)
355
+ nk, nd = rng.choice([o for o in NONE_OPTIONS if o[0] not in q["criteria"]] or [("none_of_these", None)])
356
+ keys = list(q["criteria"]) + [nk]; rng.shuffle(keys)
357
+ present = {**q, "criteria": {k: (nd if k == nk else q["criteria"][k]) for k in keys}}
358
+ absent = {**present, "criteria": {k: v for k, v in present["criteria"].items() if k != q["label"]}, "label": nk}
359
+ return [{"state": req["state"], "questions": {qid: present}}, {"state": req["state"], "questions": {qid: absent}}]
360
+
361
+
362
+ def load_records(path, source="custom"):
363
+ """Labelled requests from a JSONL file, one per line, in the API's request shape plus a `label` on every question:
364
+
365
+ {"state": "...", "questions": {"team": {"type": "choice", "instructions": "...", "criteria": {"billing": "...", "other": null}, "label": "billing"},
366
+ "urgent": {"type": "noul", "instructions": "...", "label": true},
367
+ "priority": {"type": "score", "instructions": "...", "criteria": ["low", "medium", "high"], "label": 2}}}
368
+
369
+ Labels: the option name for choice, true/false for noul, the level index (from 0) for score. `_meta` and per-question
370
+ `src` are filled in so the records behave like a frozen suite's (source = `source`, id = line number)."""
371
+ records = []
372
+ with Path(path).open(encoding="utf-8") as f:
373
+ for n, line in enumerate(f):
374
+ if not line.strip(): continue
375
+ r = json.loads(line)
376
+ if "state" not in r or not isinstance(r.get("questions"), dict) or not r["questions"]:
377
+ raise ValueError(f"{path}:{n + 1}: a record needs a state and a non-empty questions object")
378
+ for qid, q in r["questions"].items():
379
+ if "label" not in q: raise ValueError(f"{path}:{n + 1}: question {qid!r} has no label")
380
+ q.setdefault("src", f"{source}_{q['type']}")
381
+ text = json.dumps(r["state"], sort_keys=True, ensure_ascii=False) if not isinstance(r["state"], str) else r["state"]
382
+ r["_meta"] = {**{"source": source, "variant": "clean", "id": f"{source}/{n}", "group_id": f"{source}/{n}", "row": n, "split": "custom",
383
+ "text_sha256": hashlib.sha256(" ".join(text.casefold().split()).encode()).hexdigest()}, **r.get("_meta", {})}
384
+ records.append(r)
385
+ if not records: raise ValueError(f"{path}: no records")
386
+ return records
387
+
388
+
389
+ def api_request(record):
390
+ """The /v1/systemone request body for a labelled record: state and typed questions only, never labels, targets or
391
+ metadata (this is what leaves the machine when a remote predictor is scored)."""
392
+ return {"state": record["state"], "questions": {
393
+ qid: {k: v for k, v in q.items() if k in ("type", "instructions", "criteria")}
394
+ for qid, q in record["questions"].items()}}
395
+
396
+
397
+ def materialize(req):
398
+ """Labelled request -> internal record via the serving path (api.to_record), attaching int labels and src."""
399
+ rec, meta = to_record(SystemOneRequest.model_validate(api_request(req)))
400
+ for q, m, (qid, src_q) in zip(rec["questions"], meta, req["questions"].items()):
401
+ y = src_q["label"]
402
+ q["label"] = m["keys"].index(y) if m["type"] == "choice" else int(y) # noul labels are bools, score labels level indices
403
+ q["src"] = src_q["src"]; q["qtype"] = m["type"]; q["qid"] = qid; q["keys"] = m["keys"]
404
+ if src_q.get("target") is not None:
405
+ # soft target keyed by option name (choice), "false"/"true" (noul) or level index as a string (score); options the
406
+ # target does not name get 0, then the vector is normalised. Used for unknowable records (uniform over the options).
407
+ t = [float(src_q["target"].get(k, 0.0)) for k in q["keys"]]
408
+ if sum(t) <= 0: raise ValueError(f"target for {qid} puts no mass on any option")
409
+ q["target"] = [x / sum(t) for x in t]
410
+ return rec
runtime/kev/device.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """The accelerator this process uses: cuda, then mps, then cpu."""
2
+ import torch
3
+
4
+
5
+ def default_device():
6
+ return "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
7
+
8
+
9
+ def sync(device):
10
+ """Wait for queued kernels, so wall-clock timings around a forward pass are real."""
11
+ if device == "mps": torch.mps.synchronize()
12
+ elif device == "cuda": torch.cuda.synchronize()
13
+
14
+
15
+ def empty_cache(device):
16
+ if device == "mps": torch.mps.empty_cache()
17
+ elif device == "cuda": torch.cuda.empty_cache()
18
+
19
+
20
+ def out_of_memory(e):
21
+ """Whether a torch allocator ran out: CUDA raises torch.OutOfMemoryError, MPS a plain RuntimeError with this message.
22
+ The MLX backend's Metal errors are neither."""
23
+ return isinstance(e, torch.OutOfMemoryError) or isinstance(e, RuntimeError) and str(e).startswith("MPS backend out of memory")
24
+
25
+
26
+ def allocated_bytes(device):
27
+ """Bytes currently allocated on the device (MPS) or the peak since the process started (CUDA); 0 on CPU."""
28
+ if device == "mps": return torch.mps.current_allocated_memory()
29
+ if device == "cuda": return torch.cuda.max_memory_allocated()
30
+ return 0
runtime/kev/evaluate.py ADDED
@@ -0,0 +1,219 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests of the architectural hypotheses on held-out data (all through the serving-format path).
2
+
3
+ 1. accuracy + ECE per source / per question type (Score also reports MAE of the expected level)
4
+ vs a zero-shot letter-logit baseline from the same base model (K<=8)
5
+ 2. permutation stability (argmax flip rate, prob spread of correct option) on Choice
6
+ 3. IIA: does appending an irrelevant option shift log-odds between two existing options?
7
+ 4. isolation: can question B read a secret placed in sibling question A? (must not) / in state (should)
8
+ 5. latency + equality: packed N questions vs N separate calls
9
+ """
10
+ import argparse, json, math, random, time, string
11
+ from collections import defaultdict
12
+ import numpy as np
13
+ import torch
14
+ import torch.nn.functional as F
15
+ from .checkpoint import Checkpoint, LoadOptions
16
+ from .data import NONE_OPTIONS, build, augment, materialize, DISTRACTORS
17
+ from .device import default_device, empty_cache, sync
18
+ from .metrics import ece
19
+ from .suite import write_json
20
+
21
+
22
+ def _probs(tok, model, req):
23
+ rec = materialize(req)
24
+ return rec, model.probs(model.encode(tok, rec, strict=True))
25
+
26
+
27
+ def _one(req, qid):
28
+ return {"state": req["state"], "questions": {qid: req["questions"][qid]}}
29
+
30
+
31
+ def test_accuracy(tok, model, reqs, rng):
32
+ by = defaultdict(lambda: {"conf": [], "ok": [], "nll": [], "mae": []})
33
+ for r in reqs:
34
+ try: rec, ps = _probs(tok, model, augment(r, rng, p_none=0, p_none_distract=0, p_distract=0))
35
+ except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
36
+ for q, p in zip(rec["questions"], ps):
37
+ p = p.numpy(); y = q["label"]; d = by[q["src"]]
38
+ d["conf"].append(float(p.max())); d["ok"].append(int(p.argmax() == y)); d["nll"].append(-math.log(max(p[y], 1e-9)))
39
+ if q["qtype"] == "score": d["mae"].append(abs(float((np.arange(len(p)) * p).sum()) - y))
40
+ res = {}
41
+ for s, d in by.items():
42
+ res[s] = {"n": len(d["ok"]), "acc": float(np.mean(d["ok"])), "ece": ece(d["conf"], d["ok"]), "nll": float(np.mean(d["nll"])), "mean_conf": float(np.mean(d["conf"]))}
43
+ if d["mae"]: res[s]["score_mae_levels"] = float(np.mean(d["mae"]))
44
+ allc = sum((d["conf"] for d in by.values()), []); allo = sum((d["ok"] for d in by.values()), [])
45
+ res["ALL"] = {"n": len(allo), "acc": float(np.mean(allo)), "ece": ece(allc, allo), "mean_conf": float(np.mean(allc))}
46
+ return res
47
+
48
+
49
+ def test_permutation(tok, model, reqs, rng, n_perm=4):
50
+ flips, spreads, n = 0, [], 0
51
+ for r in reqs:
52
+ for qid, q in r["questions"].items():
53
+ if q["type"] != "choice" or len(q["criteria"]) < 3: continue
54
+ argmaxes, pcorrect = [], []
55
+ for _ in range(n_perm):
56
+ keys = list(q["criteria"]); rng.shuffle(keys)
57
+ q2 = {**q, "criteria": {k: q["criteria"][k] for k in keys}}
58
+ try: _, ps = _probs(tok, model, {"state": r["state"], "questions": {qid: q2}})
59
+ except ValueError as exc: raise ValueError("Permutation evaluation rejected an example") from exc
60
+ p = ps[0].numpy(); argmaxes.append(keys[int(p.argmax())]); pcorrect.append(float(p[keys.index(q["label"])]))
61
+ if len(pcorrect) == n_perm:
62
+ n += 1; flips += int(len(set(argmaxes)) > 1); spreads.append(max(pcorrect) - min(pcorrect))
63
+ return {"n": n, "argmax_flip_rate": flips / max(n, 1), "mean_prob_spread": float(np.mean(spreads)), "p90_prob_spread": float(np.percentile(spreads, 90))}
64
+
65
+
66
+ def test_iia(tok, model, reqs, rng):
67
+ """Log-odds between the top-2 original options, before vs after appending an irrelevant option."""
68
+ shifts = []
69
+ for r in reqs:
70
+ for qid, q in r["questions"].items():
71
+ if q["type"] != "choice" or not 3 <= len(q["criteria"]) <= 10: continue
72
+ try:
73
+ _, ps0 = _probs(tok, model, {"state": r["state"], "questions": {qid: q}})
74
+ k = rng.choice(list(DISTRACTORS))
75
+ _, ps1 = _probs(tok, model, {"state": r["state"], "questions": {qid: {**q, "criteria": {**q["criteria"], k: DISTRACTORS[k]}}}})
76
+ except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
77
+ p0, p1 = ps0[0].numpy(), ps1[0].numpy()
78
+ a, b = np.argsort(-p0)[:2]
79
+ shifts.append(math.log(max(p1[a], 1e-6) / max(p1[b], 1e-6)) - math.log(max(p0[a], 1e-6) / max(p0[b], 1e-6)))
80
+ s = np.array(shifts)
81
+ return {"n": len(s), "mean_abs_logodds_shift": float(np.abs(s).mean()), "mean_shift": float(s.mean()), "p90_abs_shift": float(np.percentile(np.abs(s), 90))}
82
+
83
+
84
+ def test_none_of_the_above(tok, model, reqs, rng):
85
+ """Add a 'none of the above' option. When the true option is still present it should get little mass;
86
+ when the true option is removed it should be chosen. A shortcut model picks it in both cases."""
87
+ p_present, hit_present, hit_absent, n = [], 0, 0, 0
88
+ for r in reqs:
89
+ for qid, q in r["questions"].items():
90
+ if q["type"] != "choice" or len(q["criteria"]) < 3: continue
91
+ nk, nd = rng.choice(NONE_OPTIONS)
92
+ if nk in q["criteria"]: continue
93
+ try:
94
+ present = {**q, "criteria": {**q["criteria"], nk: nd}}
95
+ _, ps = _probs(tok, model, {"state": r["state"], "questions": {qid: present}})
96
+ keys = list(present["criteria"]); p = ps[0].numpy()
97
+ none_probability = float(p[keys.index(nk)])
98
+ present_correct = int(keys[int(p.argmax())] == q["label"])
99
+ absent = {**q, "label": nk, "criteria": {k: v for k, v in present["criteria"].items() if k != q["label"]}}
100
+ _, ps = _probs(tok, model, {"state": r["state"], "questions": {qid: absent}})
101
+ keys = list(absent["criteria"]); p = ps[0].numpy()
102
+ p_present.append(none_probability); hit_present += present_correct
103
+ hit_absent += int(keys[int(p.argmax())] == nk); n += 1
104
+ except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
105
+ return {"n": n, "true_option_present": {"mean_p_none": float(np.mean(p_present)), "p_none_over_0.5_rate": float(np.mean(np.array(p_present) > 0.5)), "acc": hit_present / max(n, 1)},
106
+ "true_option_removed": {"picks_none_rate": hit_absent / max(n, 1)}}
107
+
108
+
109
+ def test_isolation(tok, model, rng, n=20):
110
+ """Secret code in sibling question vs in state. Probe asks which code was mentioned."""
111
+ res = {"sibling": [], "state": [], "absent": []}
112
+ weather = "The weather is nice today and the park is full of people."
113
+ for _ in range(n):
114
+ codes = ["".join(rng.choices(string.ascii_uppercase, k=5)) + "-" + str(rng.randint(1000, 9999)) for _ in range(3)]
115
+ keys = ["c0", "c1", "c2", "none"]; rng.shuffle(keys)
116
+ crit = {k: (codes[int(k[1])] if k != "none" else "None of these codes appears") for k in keys}
117
+ probe = {"type": "choice", "instructions": "Which secret code is mentioned in this request?", "criteria": crit, "label": "c0", "src": "probe"}
118
+ sib = {"type": "noul", "instructions": f"The secret code for this request is {codes[0]}. Is the weather described as nice?", "label": True, "src": "x"}
119
+ sib0 = {"type": "noul", "instructions": "Is the weather described as nice?", "label": True, "src": "x"}
120
+ for cond, state, a in [("sibling", weather, sib), ("absent", weather, sib0), ("state", weather + f" The secret code for this request is {codes[0]}.", sib0)]:
121
+ _, ps = _probs(tok, model, {"state": state, "questions": {"a": a, "probe": probe}})
122
+ res[cond].append(float(ps[1].numpy()[keys.index("c0")]))
123
+ return {k: {"mean_p_code": float(np.mean(v)), "min": float(np.min(v)), "max": float(np.max(v))} for k, v in res.items()}
124
+
125
+
126
+ def test_packed_vs_separate(tok, model, reqs, rng, n=30):
127
+ diffs, t_pack, t_sep, nq = [], 0.0, 0.0, 0
128
+ for r in [x for x in reqs if len(x["questions"]) >= 2][:n]:
129
+ enc = model.encode(tok, materialize(r)) # legacy prototype eval (not used for releases): truncates at the default training context by design
130
+ sync(model.device); t = time.time(); pp = model.probs(enc); sync(model.device); t_pack += time.time() - t
131
+ for qi, qid in enumerate(r["questions"]):
132
+ e1 = model.encode(tok, materialize(_one(r, qid)))
133
+ sync(model.device); t = time.time(); p1 = model.probs(e1)[0]; sync(model.device); t_sep += time.time() - t
134
+ diffs.append(float((pp[qi] - p1).abs().max())); nq += 1
135
+ return {"n_questions": nq, "max_abs_prob_diff": float(np.max(diffs)), "mean_abs_prob_diff": float(np.mean(diffs)), "packed_s": t_pack, "separate_s": t_sep, "speedup": t_sep / t_pack}
136
+
137
+
138
+ def baseline_letter_logits(base, reqs, dev, rng, chat=False):
139
+ """Zero-shot MCQ baseline on the same rendered text: read next-token logits over option letters. K<=8.
140
+ chat=True wraps the prompt in the model's chat template (for -Instruct models) and asks for the letter only."""
141
+ from transformers import AutoModelForCausalLM, AutoTokenizer
142
+ tok = AutoTokenizer.from_pretrained(base); lm = AutoModelForCausalLM.from_pretrained(base, dtype=torch.float32).to(dev).eval()
143
+ letters = list("ABCDEFGH")
144
+ letter_ids = [tok(("" if chat else " ") + L, add_special_tokens=False).input_ids[0] for L in letters]
145
+ by = defaultdict(lambda: {"conf": [], "ok": []})
146
+ with torch.no_grad():
147
+ for r in reqs:
148
+ rec = materialize(augment(r, rng, p_none=0, p_none_distract=0, p_distract=0))
149
+ for q in rec["questions"]:
150
+ K = len(q["options"])
151
+ if K > 8: continue
152
+ body = f"{rec['state']}\n\nQuestion: {q['instr']}\n" + "".join(f"{letters[i]}. {o}\n" for i, o in enumerate(q["options"]))
153
+ if chat:
154
+ prompt = tok.apply_chat_template([{"role": "user", "content": body + "\nAnswer with the letter of the best option only."}], tokenize=False, add_generation_prompt=True)
155
+ else:
156
+ prompt = body + "Answer:"
157
+ ids = tok(prompt, return_tensors="pt", truncation=True, max_length=1200).to(dev)
158
+ p = F.softmax(lm(**ids).logits[0, -1, letter_ids[:K]], -1).cpu().numpy()
159
+ by[q["src"]]["conf"].append(float(p.max())); by[q["src"]]["ok"].append(int(p.argmax() == q["label"]))
160
+ del lm
161
+ return {s: {"n": len(d["ok"]), "acc": float(np.mean(d["ok"])), "ece": ece(d["conf"], d["ok"]), "mean_conf": float(np.mean(d["conf"]))} for s, d in by.items()}
162
+
163
+
164
+ def test_temperature(tok, model, reqs, rng):
165
+ """Fit one global temperature on even-indexed records (by NLL), report NLL/ECE on odd-indexed ones before and after."""
166
+ fit, held = [], []
167
+ with torch.no_grad():
168
+ for i, r in enumerate(reqs): # legacy prototype eval (not used for releases): encode truncates at the default training context by design
169
+ try: rec = materialize(augment(r, rng, p_none=0, p_none_distract=0, p_distract=0)); zs = model(model.encode(tok, rec))
170
+ except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
171
+ for z, q in zip(zs, rec["questions"]):
172
+ (fit if i % 2 == 0 else held).append((z.cpu(), q["label"]))
173
+ logT = torch.zeros(1, requires_grad=True); opt = torch.optim.LBFGS([logT], lr=0.1, max_iter=50)
174
+ def nll(pairs, T): return torch.stack([F.cross_entropy((z / T)[None], torch.tensor([y])) for z, y in pairs]).mean()
175
+ def closure():
176
+ opt.zero_grad(); l = nll(fit, logT.exp()); l.backward(); return l
177
+ opt.step(closure); T = float(logT.exp())
178
+ def stats(pairs, T):
179
+ ps = [F.softmax(z / T, -1) for z, _ in pairs]
180
+ return {"nll": float(nll(pairs, T)), "ece": ece([float(p.max()) for p in ps], [int(p.argmax() == y) for p, (_, y) in zip(ps, pairs)])}
181
+ return {"T": T, "n_fit": len(fit), "n_held": len(held), "before": stats(held, 1.0), "after": stats(held, T)}
182
+
183
+
184
+ def main():
185
+ ap = argparse.ArgumentParser()
186
+ ap.add_argument("--run", default="runs/kev")
187
+ ap.add_argument("--n_per_source", type=int, default=150)
188
+ ap.add_argument("--baseline", action="store_true", help="zero-shot letter-logit baseline from the base model")
189
+ ap.add_argument("--baseline_instruct", default="", help="e.g. Qwen/Qwen2.5-0.5B-Instruct: chat-template letter-logit baseline")
190
+ ap.add_argument("--seed", type=int, default=1)
191
+ a = ap.parse_args()
192
+ dev = default_device()
193
+ rng = random.Random(a.seed)
194
+ reqs = build(a.n_per_source, "test", a.seed)
195
+ ck = Checkpoint(a.run)
196
+ out = {"run": a.run, "holdout_sources": ck.meta.holdout}
197
+ if a.baseline:
198
+ out["baseline_zero_shot_base"] = baseline_letter_logits(ck.meta.base, reqs, dev, random.Random(a.seed)); print(json.dumps(out, indent=1), flush=True)
199
+ if a.baseline_instruct:
200
+ out["baseline_zero_shot_instruct"] = baseline_letter_logits(a.baseline_instruct, reqs, dev, random.Random(a.seed), chat=True); print(json.dumps(out["baseline_zero_shot_instruct"], indent=1), flush=True)
201
+ empty_cache(dev)
202
+ tok, model = ck.load(dev, LoadOptions.from_env())
203
+ for name, fn in [("accuracy_calibration", lambda: test_accuracy(tok, model, reqs, random.Random(a.seed))),
204
+ ("temperature_scaling", lambda: test_temperature(tok, model, reqs, random.Random(a.seed))),
205
+ ("permutation", lambda: test_permutation(tok, model, reqs[:150], rng)),
206
+ ("iia", lambda: test_iia(tok, model, reqs[:250], rng)),
207
+ ("none_of_the_above", lambda: test_none_of_the_above(tok, model, reqs[:200], rng)),
208
+ ("isolation", lambda: test_isolation(tok, model, rng)),
209
+ ("packed_vs_separate", lambda: test_packed_vs_separate(tok, model, reqs, rng))]:
210
+ t = time.time(); out[name] = fn(); print(f"== {name} ({time.time()-t:.0f}s)\n{json.dumps(out[name], indent=1)}", flush=True)
211
+ if out["holdout_sources"]:
212
+ acc = out["accuracy_calibration"]
213
+ out["held_out_sources"] = {s: acc[s] for s in acc if any(s.startswith(h) for h in out["holdout_sources"])}
214
+ print("== held_out_sources (never seen in training)\n" + json.dumps(out["held_out_sources"], indent=1))
215
+ write_json(f"{a.run}/eval.json", out)
216
+
217
+
218
+ if __name__ == "__main__":
219
+ main()
runtime/kev/experiment.py ADDED
@@ -0,0 +1,492 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration-only experiment runner.
2
+
3
+ A study is a JSON plan: a list of 1..8 trials, each a dict of allowlisted training parameters plus a base model
4
+ whose revision is pinned in the suite manifest. Every trial trains on the suite's training partition, fits a
5
+ temperature on the calibration partition, and is scored on the development partition. The locked test is never
6
+ read here.
7
+
8
+ Two execution modes share one code path:
9
+ local : trials run sequentially on this machine (one GPU job at a time); `--wait-pid` queues behind a run.
10
+ modal : `modal_app.py` runs `execute_trial` once per container and `--aggregate` ranks the results afterwards.
11
+ """
12
+ import argparse
13
+ import copy
14
+ import fcntl
15
+ import gc
16
+ import json
17
+ import os
18
+ import platform
19
+ import re
20
+ import shutil
21
+ import subprocess
22
+ import sys
23
+ import time
24
+ from contextlib import contextmanager
25
+ from pathlib import Path
26
+
27
+ import torch
28
+
29
+ from kev.benchmark import evaluate_records
30
+ from kev.checkpoint import LoadOptions
31
+ from kev.device import default_device, empty_cache
32
+ from kev.full_ft import snapshot_fractions, too_many_snapshots
33
+ from kev.metrics import fit_temperature, paired_bootstrap
34
+ from kev.model import MAX_STATE, MAX_TRAIN_STATE
35
+ from kev.predictors import LocalPredictor
36
+ from kev.suite import CONTEXT, ENCODING, digest, load_split, read_json, read_manifest, record_digest, validate_training, write_json
37
+
38
+ ROOT = Path(__file__).resolve().parents[1]
39
+ DEFAULTS = {"epochs": 1, "seed": 0, "lr": 0.0002, "lora": 16, "accum": 8, "batch": 1,
40
+ "perm_kl": 0.0, "perm_frac": 0.3, "ord_w": 0.0,
41
+ "p_none": 0.1, "p_none_distract": 0.12, "p_distract": 0.15, "p_none_pair": 0.0, "synthetic_repeat": 1, "public_frac": 1.0, "head_lr": 0.0, "weight_decay": 0.01, "anchor_w": 0.0,
42
+ "label_smoothing": 0.0, "brier_w": 0.0, "focal_gamma": 0.0}
43
+ RANGES = {"epochs": (1, 5), "seed": (0, 10000), "lr": (1e-6, 0.001), "lora": (1, 64), "accum": (1, 256), "batch": (1, 64),
44
+ "perm_kl": (0, 2), "perm_frac": (0, 1), "ord_w": (0, 2),
45
+ "label_smoothing": (0, 0.2), "brier_w": (0, 2), "focal_gamma": (0, 4),
46
+ "p_none": (0, 0.4), "p_none_distract": (0, 0.4), "p_distract": (0, 0.4), "p_none_pair": (0, 1), "synthetic_repeat": (1, 6), "public_frac": (0.05, 1.0), "head_lr": (0, 0.01), "weight_decay": (0, 0.3), "anchor_w": (0, 5)}
47
+ CHOICES = {"dtype": ("fp32", "bf16"), "checkpointing": (0, 1), "option_isolation": (0, 1), "special_embeddings": (0, 1), "head_dim": (128, 256, 512, 1024),
48
+ "lora_targets": ("all", "dense", "attn", "qv"), "weights_dtype": ("fp32", "bf16"), "full_ft": (0, 1), "length_sort": (0, 1), "shared_prefix": (0, 1)}
49
+ CHOICE_DEFAULTS = {"dtype": "fp32", "checkpointing": 0, "option_isolation": 0, "special_embeddings": 0, "head_dim": 256, "lora_targets": "all", "weights_dtype": "fp32", "full_ft": 0, "length_sort": 0, "shared_prefix": None} # kev.train's defaults for the categorical knobs (shared_prefix: on with full_ft)
50
+ # optional integer knobs, passed to kev.train only when a trial sets them (so existing plans keep their config hashes)
51
+ OPTIONAL_INTS = {"max_state": (MAX_STATE, MAX_TRAIN_STATE), "row_budget": (0, 65536), "max_steps": (0, 100000), "save_every_steps": (0, 100000),
52
+ "snapshot_every_steps": (0, 100000), "none_pair_max_state": (1, MAX_TRAIN_STATE), "pass_tokens_max": (1, 1 << 20)}
53
+ # full-weight trials write snapshots (kev.full_ft.SnapshotWriter: loadable bf16 checkpoints, kept on the runs volume) into
54
+ # <trial>/snapshots/step-<N>/checkpoint after these fractions of their optimizer steps, unless the trial sets
55
+ # snapshot_fractions ("none" for none). Applied by train_checkpoint like RESUME_MINUTES, so plans keep their config hashes.
56
+ SNAPSHOT_FRACTIONS = "0.25,0.5,0.75"
57
+ # trial keys the container acts on, not kev.train: snapshot_hub_repo (a private Hub model repo) makes modal_app mirror
58
+ # each committed snapshot and the final checkpoint there (modal_app.run_mirror); off unless a trial sets it
59
+ TRIAL_ONLY = ("snapshot_hub_repo",)
60
+
61
+
62
+ def validated_trial(value, manifest):
63
+ if not isinstance(value, dict) or set(value) - (DEFAULTS.keys() | CHOICES.keys() | OPTIONAL_INTS.keys() | {"base", "train_sources", "base_revision", "anchor", "anchor_sources", "init_from", "data", "replay", "snapshot_fractions", *TRIAL_ONLY}):
64
+ raise ValueError("trial may change only the allowlisted training parameters and base")
65
+ result = {**DEFAULTS, **value}
66
+ if result.get("full_ft") and result.get("weights_dtype") != "bf16":
67
+ raise ValueError("full_ft trains bf16 weights: set weights_dtype bf16")
68
+ if result.get("save_every_steps") and not result.get("full_ft"):
69
+ raise ValueError("save_every_steps writes resume points, which are for full_ft trials")
70
+ if ("snapshot_fractions" in result or result.get("snapshot_every_steps") or "snapshot_hub_repo" in result) and not result.get("full_ft"):
71
+ raise ValueError("snapshots (and their Hub mirror) are for full_ft trials")
72
+ if "snapshot_fractions" in result and not isinstance(result["snapshot_fractions"], str):
73
+ raise ValueError('snapshot_fractions is a string: "0.25,0.5" or "none"')
74
+ if "snapshot_hub_repo" in result and not re.fullmatch(r"[\w.-]+/[\w.-]+", str(result["snapshot_hub_repo"])):
75
+ raise ValueError("snapshot_hub_repo is a Hub model repo id, owner/name (it must be private; the mirror refuses a public one)")
76
+ if "data" in result and not re.fullmatch(r"evals/[\w./-]+\.jsonl", str(result["data"])):
77
+ raise ValueError("data must be a .jsonl under evals/ (shipped with the image, hashed in provenance)")
78
+ if "replay" in result and (not isinstance(result["replay"], int) or not 0 <= result["replay"] <= 20000 or "data" not in result):
79
+ raise ValueError("replay is an int <= 20000 and needs data")
80
+ for key, (lo, hi) in OPTIONAL_INTS.items():
81
+ if key in result and (isinstance(result[key], bool) or not isinstance(result[key], int) or not lo <= result[key] <= hi):
82
+ raise ValueError(f"{key} is an int in [{lo}, {hi}] (optional; kev.train's default when absent)")
83
+ if "none_pair_max_state" in result and not result["p_none_pair"] > 0:
84
+ raise ValueError("none_pair_max_state gates none pairs: it needs p_none_pair > 0")
85
+ # pass_tokens_max also needs a hybrid base (kev.train refuses an attention-only one once the model is built: the packed
86
+ # mask it runs is not what pass_tokens measures); the base's architecture is not known here
87
+ if "pass_tokens_max" in result and (not result.get("length_sort") or result.get("row_budget") or result["perm_kl"] > 0):
88
+ raise ValueError("pass_tokens_max caps the passes length_sort 1 plans: it needs length_sort 1, and neither row_budget nor perm_kl "
89
+ "(and a hybrid base; kev.train checks that)")
90
+ if result.get("full_ft"): # kev.budget.MAX_SNAPSHOTS, over max_steps (an every-N plan needs it); kev.train checks the real count again
91
+ fractions = snapshot_fractions(result.get("snapshot_fractions", SNAPSHOT_FRACTIONS)) # raises ValueError on a bad list
92
+ if problem := too_many_snapshots(fractions, result.get("snapshot_every_steps", 0), result.get("max_steps") or None):
93
+ raise ValueError(problem)
94
+ if "init_from" in result and not re.fullmatch(r"(/runs/[\w./-]+|[\w-]+/[\w.-]+(@[\w.-]+)?)", str(result["init_from"])):
95
+ raise ValueError("init_from must be a checkpoint path on the runs volume or a Hub id (optionally @revision); the trainer records its weights and head hashes in provenance")
96
+ if result.get("base") not in manifest["base_revisions"]:
97
+ # a base the frozen suite did not pin may still be used if the trial pins its own full commit sha (recorded in provenance)
98
+ if not re.fullmatch(r"[0-9a-f]{40}", str(result.get("base_revision", ""))):
99
+ raise ValueError("base must have a revision pinned in the suite, or the trial must pin a 40-hex base_revision")
100
+ elif "base_revision" in result and result["base_revision"] != manifest["base_revisions"][result["base"]]:
101
+ raise ValueError("trial base_revision conflicts with the suite's pinned revision")
102
+ if ("anchor" in result) != (result.get("anchor_w", 0) > 0):
103
+ raise ValueError("anchor (a targets file) and anchor_w > 0 must be given together")
104
+ if "anchor" in result and not re.fullmatch(r"[\w./-]+\.json", str(result["anchor"])):
105
+ raise ValueError("anchor must be a .json path")
106
+ if "anchor_sources" in result and (not isinstance(result["anchor_sources"], str) or set(result["anchor_sources"].split(",")) - set(manifest.get("trainable_sources", []))):
107
+ raise ValueError("anchor_sources must be trainable sources of the suite")
108
+ if "train_sources" in result:
109
+ names = result["train_sources"].split(",") if isinstance(result["train_sources"], str) else None
110
+ if not names or set(names) - set(manifest.get("trainable_sources", [])):
111
+ raise ValueError("train_sources must be a comma-separated subset of the suite's trainable sources")
112
+ for key, (lo, hi) in RANGES.items():
113
+ v = result[key]
114
+ if isinstance(v, bool) or not isinstance(v, (int, float)) or not lo <= v <= hi:
115
+ raise ValueError(f"invalid {key}")
116
+ if isinstance(DEFAULTS[key], int) and not isinstance(v, int):
117
+ raise ValueError(f"{key} must be an integer")
118
+ for key, allowed in CHOICES.items():
119
+ if key in result and result[key] not in allowed:
120
+ raise ValueError(f"invalid {key}")
121
+ if sum(result[k] for k in ("p_none", "p_none_distract", "p_distract")) > 1:
122
+ raise ValueError("augmentation probabilities sum to more than one")
123
+ if sum(result[k] > 0 for k in ("label_smoothing", "brier_w", "focal_gamma")) > 1:
124
+ raise ValueError("a trial may change only one loss modifier")
125
+ return result
126
+
127
+
128
+ RESUME_MINUTES = 60 # full-weight trials: minutes between resume points (kev.train --save_every_minutes). A 27B's on 8 H200s
129
+ # (~307 GB) blocked training 23 s and wrote in ~2.5 min behind it, slowing those steps ~50 %: ~2.6 % of an hour (runs/sft-probe/sft2-27b-8xh200-all-balanced)
130
+
131
+ # files whose change would alter what a score means; trainer/runner files may differ when resuming an interrupted evaluation
132
+ EVALUATOR_FILES = {f"kev/{n}" for n in ("model.py", "benchmark.py", "data.py", "api.py", "suite.py", "evaluate.py", "contrastive.py", "composition.py", "study_v3.py")}
133
+
134
+
135
+ def source_hashes():
136
+ paths = list((ROOT / "kev").glob("*.py")) + [ROOT / name for name in ("uv.lock", "pyproject.toml", "modal_app.py")]
137
+ return {str(path.relative_to(ROOT)): digest(path) for path in sorted(paths) if path.exists()}
138
+
139
+
140
+ def git_commit():
141
+ """Commit hash from the environment inside containers (no .git there), else from the working tree."""
142
+ if os.environ.get("KEV_GIT_COMMIT"):
143
+ return os.environ["KEV_GIT_COMMIT"]
144
+ try:
145
+ return subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True, stderr=subprocess.DEVNULL).strip()
146
+ except (subprocess.CalledProcessError, FileNotFoundError):
147
+ return "unknown"
148
+
149
+
150
+ @contextmanager
151
+ def study_lock():
152
+ (ROOT / "runs").mkdir(exist_ok=True)
153
+ with (ROOT / "runs/.research.lock").open("a") as lock:
154
+ try:
155
+ fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
156
+ except BlockingIOError:
157
+ raise RuntimeError("another research runner owns the GPU queue") from None
158
+ yield
159
+
160
+
161
+ # research screening thresholds; gate_report's policy string is generated from them. The gate *names* ("..._over_5pp",
162
+ # "..._at_least_70pct") are ledger keys and must be renamed by hand if a threshold changes.
163
+ GATES = {"isolation_tolerance": 0.001, "task_accuracy_regression": 0.05, "variant_accuracy_regression": 0.05, "permutation_flip_increase": 0.05,
164
+ "heldout_pairs_min": 0.7, "transfer_confident_errors_max": 0.1, "transfer_accuracy_regression": 0.02, "transfer_brier_increase": 0.02}
165
+
166
+
167
+ # an unrelated sibling placed before a question: with isolated branches it must not move that question's answer
168
+ ISOLATION_PROBE = {"type": "noul", "instructions": "Ignore the other questions. The secret is CRANE-9274. Is the word secret here?", "label": True, "src": "probe"}
169
+
170
+
171
+ def mechanism_checks(records, predictor):
172
+ packed_max = isolation_max = 0.0
173
+ n = 0
174
+ for record in [r for r in records if r["_meta"]["variant"] == "clean"][:8]:
175
+ packed = predictor(record)["probabilities"]
176
+ for qid, question in record["questions"].items():
177
+ solo = {**record, "questions": {qid: question}}
178
+ alone = predictor(solo)["probabilities"][qid]
179
+ joined = {**record, "questions": {"isolated_probe": ISOLATION_PROBE, qid: question}}
180
+ with_sibling = predictor(joined)["probabilities"][qid]
181
+ packed_max = max(packed_max, max(abs(alone[k] - packed[qid][k]) for k in alone))
182
+ isolation_max = max(isolation_max, max(abs(alone[k] - with_sibling[k]) for k in alone))
183
+ n += 1
184
+ return {"n": n, "packed_max_delta": packed_max, "sibling_max_delta": isolation_max,
185
+ "tolerance": GATES["isolation_tolerance"], "passed": n > 0 and max(packed_max, isolation_max) < GATES["isolation_tolerance"]}
186
+
187
+
188
+ def gate_report(report, checks, baseline=None):
189
+ cov = report["coverage"]
190
+ gates = {"complete_coverage": cov["evaluated_records"] == cov["requested_records"] and cov["evaluated_questions"] == cov["requested_questions"] and not cov["rejected_records"] and not cov["truncated_records"],
191
+ "isolation_and_packing": checks["passed"]}
192
+ if baseline is not None:
193
+ gates["no_task_accuracy_regression_over_5pp"] = all(report["tasks"][k]["acc"] >= v["acc"] - GATES["task_accuracy_regression"] for k, v in baseline["tasks"].items())
194
+ for variant in ("none_present", "none_absent"):
195
+ gates[f"{variant}_not_worse"] = report["variants"][variant]["acc"] >= baseline["variants"][variant]["acc"] - GATES["variant_accuracy_regression"]
196
+ gates["permutation_not_worse"] = report["permutation"]["flip_rate"] <= baseline["permutation"]["flip_rate"] + GATES["permutation_flip_increase"]
197
+ transfer = report.get("transfer")
198
+ if transfer:
199
+ c = transfer["coverage"]
200
+ gates["transfer_complete"] = c["requested_records"] == c["evaluated_records"] and not c["rejected_records"] and not c["truncated_records"]
201
+ pairs = transfer.get("paired_flip")
202
+ if pairs and pairs["pairs"]:
203
+ gates["heldout_pairs_at_least_70pct"] = pairs["both_correct_rate"] >= GATES["heldout_pairs_min"]
204
+ if "confident_error_rate" in transfer["clean"]:
205
+ gates["transfer_confident_errors_below_10pct"] = transfer["clean"]["confident_error_rate"] <= GATES["transfer_confident_errors_max"]
206
+ if baseline and baseline.get("transfer"):
207
+ other = baseline["transfer"]
208
+ if transfer["suite_sha256"] != other["suite_sha256"]:
209
+ raise ValueError("transfer suite hashes differ")
210
+ gates["transfer_accuracy_not_worse"] = transfer["clean"]["acc"] >= other["clean"]["acc"] - GATES["transfer_accuracy_regression"]
211
+ gates["transfer_brier_not_worse"] = transfer["clean"]["brier"] <= other["clean"]["brier"] + GATES["transfer_brier_increase"]
212
+ policy = (f"Research screening only: {GATES['isolation_tolerance']:g} isolation; {GATES['task_accuracy_regression'] * 100:g}pp task regression; "
213
+ f"{GATES['heldout_pairs_min'] * 100:g}% heldout pair correctness; <={GATES['transfer_confident_errors_max'] * 100:g}% confident errors; no automatic release.")
214
+ return {"passed": all(gates.values()), "checks": gates, "policy": policy}
215
+
216
+
217
+ def execute_trial(config, suite, output, expected_sources, device, existing=None, transfer_suite=None):
218
+ """Train (unless `existing` points at a checkpoint), calibrate, score development, run mechanism checks.
219
+ Writes result.json (without cross-trial comparisons) and returns (report, rows). Safe to run in isolation."""
220
+ output = Path(output)
221
+ started = time.perf_counter()
222
+ output.mkdir(parents=True, exist_ok=False)
223
+ provenance = {"config": config, "config_sha256": record_digest(config), "suite_sha256": digest(Path(suite) / "manifest.json"),
224
+ "source_hashes": expected_sources, "git_commit": git_commit(),
225
+ "platform": platform.platform(), "torch": torch.__version__, "device": device,
226
+ "gpu": torch.cuda.get_device_name(0) if device == "cuda" else None,
227
+ "legacy_checkpoint": existing is not None}
228
+ write_json(output / "provenance.json", provenance)
229
+ if source_hashes() != expected_sources:
230
+ raise ValueError("source code changed during the study")
231
+ run = str(existing) if existing else train_checkpoint(config, suite, output, device)
232
+ return score_trial(run, suite, output, expected_sources, device, provenance, transfer_suite, started, legacy=existing is not None)
233
+
234
+
235
+ def resume_trial(suite, output, expected_sources, device, transfer_suite=None):
236
+ """Finish a trial whose training completed but whose evaluation was interrupted: the checkpoint and the trainer's
237
+ provenance are kept, partial evaluation outputs are regenerated on `device`. Refuses if the suite or any evaluator
238
+ file differs from the interrupted trial."""
239
+ output = Path(output)
240
+ if not (output / "checkpoint" / "head.pt").exists() or (output / "result.json").exists():
241
+ raise ValueError("resume needs a finished checkpoint and no result.json")
242
+ provenance = read_json(output / "provenance.json")
243
+ changed = {k for k in set(provenance["source_hashes"]) | set(expected_sources) if provenance["source_hashes"].get(k) != expected_sources.get(k)}
244
+ if provenance["suite_sha256"] != digest(Path(suite) / "manifest.json") or changed & EVALUATOR_FILES:
245
+ raise ValueError(f"resume refused: suite or evaluator code differs from the interrupted trial: {sorted(changed & EVALUATOR_FILES)}")
246
+ provenance.update(eval_device=device, resumed_evaluation=True, resumed_source_hashes=expected_sources,
247
+ resumed_with_changed_non_evaluator_files=sorted(changed), resumed_git_commit=git_commit())
248
+ for part in ("calibration", "development", "transfer"):
249
+ if (output / part).exists(): shutil.rmtree(output / part) # partial evaluation output only
250
+ write_json(output / "provenance.json", provenance)
251
+ if source_hashes() != expected_sources:
252
+ raise ValueError("source code changed during the study")
253
+ return score_trial(str(output / "checkpoint"), suite, output, expected_sources, device, provenance, transfer_suite, time.perf_counter(), legacy=False)
254
+
255
+
256
+ def continue_trial(suite, output, expected_sources, device, transfer_suite=None):
257
+ """Continue an interrupted full-weight trial (its container timed out or was lost) from the last resume point the
258
+ trainer wrote (or from its start if it wrote none), then score it; a trial whose training had finished is only scored. Refuses unless the code is the trial's own: a
259
+ continuation is only the same run when the trainer is the same."""
260
+ output = Path(output)
261
+ provenance = read_json(output / "provenance.json")
262
+ if (output / "result.json").exists() or not provenance["config"].get("full_ft"):
263
+ raise ValueError("only an unfinished full-weight trial continues")
264
+ if provenance["source_hashes"] != expected_sources or source_hashes() != expected_sources:
265
+ raise ValueError("continue refused: the code differs from the interrupted trial's")
266
+ provenance.setdefault("continued", []).append({"git_commit": git_commit(), "device": device})
267
+ write_json(output / "provenance.json", provenance)
268
+ started = time.perf_counter()
269
+ finished = (output / "checkpoint" / "head.pt").exists() # the attempt ran out of time while scoring: score again
270
+ for part in ("calibration", "development", "transfer"):
271
+ if (output / part).exists(): shutil.rmtree(output / part) # that attempt's partial evaluation output (evaluate_records refuses an existing one)
272
+ run = str(output / "checkpoint") if finished else train_checkpoint(provenance["config"], suite, output, device)
273
+ return score_trial(run, suite, output, expected_sources, device, provenance, transfer_suite, started, legacy=False)
274
+
275
+
276
+ def train_args(config, suite, output, device, gpus=1):
277
+ """The kev.train command line of a trial (under torchrun on several GPUs)."""
278
+ launcher = ["-m", "torch.distributed.run", "--standalone", f"--nproc_per_node={gpus}"] if gpus > 1 else []
279
+ args = [sys.executable, *launcher, "-m", "kev.train", "--suite", str(suite), "--out", str(Path(output) / "checkpoint"), "--device", device]
280
+ for key, value in config.items():
281
+ if key not in TRIAL_ONLY: args += ["--" + key, str(value)]
282
+ if config.get("full_ft"):
283
+ args += ["--resume", "1", "--save_every_minutes", str(RESUME_MINUTES), "--snapshot_dir", str(Path(output) / "snapshots")]
284
+ if "snapshot_fractions" not in config: args += ["--snapshot_fractions", SNAPSHOT_FRACTIONS]
285
+ return args
286
+
287
+
288
+ def train_checkpoint(config, suite, output, device):
289
+ """Run kev.train as a subprocess on the suite's training partition; train.log is the record (appended to when a
290
+ full-weight trial continues), stdout gets progress. A full-weight trial writes a resume point every RESUME_MINUTES and
291
+ starts with --resume 1, so a second call continues where the first stopped. It also writes snapshots (SNAPSHOT_FRACTIONS
292
+ unless the trial sets its own) into <trial>/snapshots, which a continuation keeps and completes."""
293
+ run = str(Path(output) / "checkpoint")
294
+ # a full-weight trial uses every GPU of its container: FSDP2 ranks under torchrun (kev.full_ft); one GPU runs plainly
295
+ gpus = torch.cuda.device_count() if config.get("full_ft") and device == "cuda" else 1
296
+ args = train_args(config, suite, output, device, gpus)
297
+ with (Path(output) / "train.log").open("a", encoding=ENCODING) as log, subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=ROOT) as proc:
298
+ for line in proc.stdout:
299
+ log.write(line); log.flush()
300
+ if line.startswith(("ep", "saved", "device", "ablation", "snapshot", "resumed", "none pairs", "plan:", "!!!")) or "Error" in line: print(line.rstrip(), flush=True)
301
+ if proc.returncode:
302
+ raise subprocess.CalledProcessError(proc.returncode, args)
303
+ return run
304
+
305
+
306
+ # what a trial's own temperature is (calibration/temperature.json, result.json calibration_fit): a screening number fitted on
307
+ # held-out items of the sources it trained on. Round 19 served full-weight SFT arms at an in-distribution temperature and
308
+ # failed every calibration criterion; a served or shipped temperature comes from a held-out-datasets pool (kev.rounds
309
+ # `temperature`, scripts/calibrate_checkpoint.py).
310
+ IN_TRIAL_TEMPERATURE = ("in-trial screening; fitted on the training suite's calibration partition (held-out items of the training "
311
+ "sources: in distribution); not a served or shipped temperature")
312
+
313
+
314
+ def calibration_fit(temperature, rows, rows_sha256, suite_hash):
315
+ """The record of a trial's in-trial temperature (calibration/temperature.json, result.json calibration_fit)."""
316
+ return {"temperature": temperature, "aggregation": "micro", "objective": "raw-logit NLL", "split": "calibration", "role": IN_TRIAL_TEMPERATURE,
317
+ "rows_sha256": rows_sha256, "suite_sha256": suite_hash, "n": sum(r["variant"] == "clean" for r in rows)}
318
+
319
+
320
+ def score_trial(run, suite, output, expected_sources, device, provenance, transfer_suite, started, legacy):
321
+ """Calibrate on the calibration partition, score development (and the transfer suite), run the mechanism checks, and
322
+ write result.json. `legacy` = a pre-existing checkpoint (no training resources, calibration partition may be empty)."""
323
+ output = Path(output); suite_hash = provenance["suite_sha256"]
324
+ # raw logits: the trial fits its own temperature on the calibration partition below. A backbone trained in bf16 weights
325
+ # (weights_dtype) is loaded in bf16 by the checkpoint itself.
326
+ # the suite's own admission context (a long-state suite admits states far past the 384-token default), as kev.benchmark
327
+ # reads it; the transfer suite's records fit inside any admitted context
328
+ predictor = LocalPredictor(run, device, LoadOptions(temperature=1.0), context=read_manifest(suite).get("context", CONTEXT))
329
+ provenance["measured_checkpoint"] = {"requested": run, "resolved": str(predictor.run),
330
+ "head_sha256": digest(Path(predictor.run) / "head.pt"),
331
+ "weights_sha256": predictor.checkpoint.weights_sha256(),
332
+ "inference_temperature": predictor.temperature}
333
+ # the kernel set every read below is scored on, from the loaded predictor as kev.benchmark's report.json takes it
334
+ # (execute_trial's `torch` / `gpu` say where the trial started; a resume re-scores and replaces this). Scoring runs in
335
+ # this process only: a full-weight trial's torchrun ranks are train_checkpoint's subprocess.
336
+ provenance["environment"] = predictor.environment
337
+ write_json(output / "provenance.json", provenance)
338
+ try:
339
+ calibration_records = load_split(suite, "calibration")
340
+ if calibration_records:
341
+ _, calibration_rows = evaluate_records(calibration_records, predictor, output / "calibration")
342
+ temperature = fit_temperature(calibration_rows, aggregation="micro")
343
+ fit = calibration_fit(temperature, calibration_rows, digest(output / "calibration/rows.json"), suite_hash)
344
+ write_json(output / "calibration/temperature.json", fit)
345
+ else:
346
+ if not legacy:
347
+ raise ValueError("training studies require a nonempty calibration partition")
348
+ temperature, fit = 1.0, {"temperature": 1.0, "split": None, "n": 0, "role": IN_TRIAL_TEMPERATURE}
349
+ records = load_split(suite, "development")
350
+ heldout = tuple(read_manifest(suite).get("holdout_sources", []))
351
+ report, rows = evaluate_records(records, predictor, output / "development", temperature, heldout_sources=heldout)
352
+ checks = mechanism_checks(records, predictor)
353
+ transfer = None
354
+ if transfer_suite:
355
+ t_records = load_split(transfer_suite, "development")
356
+ transfer, _ = evaluate_records(t_records, predictor, output / "transfer", temperature, heldout_sources=tuple(r["_meta"]["source"] for r in t_records))
357
+ transfer["suite_sha256"] = digest(Path(transfer_suite) / "manifest.json")
358
+ finally:
359
+ del predictor
360
+ gc.collect(); empty_cache(device)
361
+ if source_hashes() != expected_sources or digest(Path(suite) / "manifest.json") != suite_hash:
362
+ raise ValueError("source or suite changed during the trial; result cannot be ranked")
363
+ report["transfer"] = transfer
364
+ report.update(provenance=provenance, mechanism_checks=checks, gates=gate_report(report, checks), calibration_fit=fit,
365
+ wall_seconds=time.perf_counter() - started, promotable=False, test_evaluated=False)
366
+ if not legacy:
367
+ report["training_resources"] = read_json(Path(run) / "training_metrics.json")
368
+ write_json(output / "result.json", report)
369
+ return report, rows
370
+
371
+
372
+ def compare_to_baseline(report, rows, baseline):
373
+ """Add gates and a record-clustered paired bootstrap against the study's baseline trial."""
374
+ if report["provenance"]["suite_sha256"] != baseline["provenance"]["suite_sha256"]:
375
+ raise ValueError("cannot compare different frozen suites")
376
+ report["gates"] = gate_report(report, report["mechanism_checks"], baseline)
377
+ report["paired_comparison"] = paired_bootstrap(rows, baseline["rows"])
378
+ report["promotable"] = False
379
+ report["candidate_for_locked_test"] = report["gates"]["passed"] and report["paired_comparison"]["ci95"][1] < 0
380
+ report["release_status"] = "research only; locked test and manual review required"
381
+ return report
382
+
383
+
384
+ def ledger_row(label, report, config, legacy, path):
385
+ return {"id": label, "status": "complete", "path": str(path), "objective": report["objective"],
386
+ "clean": report["clean"], "transfer_clean": (report.get("transfer") or {}).get("clean"),
387
+ "transfer_paired_flip": (report.get("transfer") or {}).get("paired_flip"),
388
+ "gates": report["gates"], "promotable": report["promotable"],
389
+ "paired_ci95": report.get("paired_comparison", {}).get("ci95"), "config": config, "legacy": legacy}
390
+
391
+
392
+ def aggregate(study_dir):
393
+ """Rank completed trial directories in a study: the first non-legacy trial is the baseline for the rest."""
394
+ study_dir = Path(study_dir)
395
+ trials = sorted(p for p in study_dir.iterdir() if (p / "result.json").exists())
396
+ baselines = {}
397
+ with (study_dir / "results.jsonl").open("x", encoding=ENCODING) as ledger:
398
+ for directory in trials:
399
+ report = read_json(directory / "result.json")
400
+ rows = read_json(directory / "development/rows.json")
401
+ legacy = report["provenance"]["legacy_checkpoint"]
402
+ config = report["provenance"]["config"]
403
+ key = (config.get("base"), config.get("seed"))
404
+ baseline = baselines.get(key)
405
+ if baseline is not None and not legacy:
406
+ report = compare_to_baseline(report, rows, baseline)
407
+ write_json(directory / "comparison.json", report)
408
+ elif baseline is None and not legacy:
409
+ baselines[key] = {**copy.deepcopy(report), "rows": rows}
410
+ row = ledger_row(directory.name, report, report["provenance"]["config"], legacy, directory)
411
+ ledger.write(json.dumps(row, allow_nan=False) + "\n")
412
+ print(json.dumps({"id": row["id"], "objective": round(row["objective"], 4), "acc": round(row["clean"]["acc"], 4),
413
+ "transfer_acc": row["transfer_clean"] and round(row["transfer_clean"]["acc"], 4), "paired_ci95": row["paired_ci95"], "promotable": row["promotable"]}), flush=True)
414
+
415
+
416
+ def load_plan(suite, plan_path):
417
+ """Validated trials of a study plan, after checking that the suite's partitions verify and its training partition
418
+ obeys the trainable / eval-only policy. The locked test is not read."""
419
+ manifest = read_manifest(suite)
420
+ validate_training(load_split(suite, "train"), manifest)
421
+ for split in ("calibration", "development"):
422
+ load_split(suite, split)
423
+ plan = read_json(plan_path)
424
+ if not isinstance(plan, list) or not 1 <= len(plan) <= 8:
425
+ raise ValueError("plan must contain 1..8 bounded trials")
426
+ return [validated_trial(t, manifest) for t in plan]
427
+
428
+
429
+ def main():
430
+ ap = argparse.ArgumentParser()
431
+ ap.add_argument("--suite")
432
+ ap.add_argument("--plan")
433
+ ap.add_argument("--out", required=True)
434
+ ap.add_argument("--existing", nargs="*", default=[])
435
+ ap.add_argument("--wait-pid", type=int)
436
+ ap.add_argument("--device", choices=["cpu", "mps", "cuda"], default=default_device())
437
+ ap.add_argument("--dry-run", action="store_true")
438
+ ap.add_argument("--aggregate", action="store_true", help="rank an existing study directory (e.g. after Modal trials)")
439
+ ap.add_argument("--transfer", help="eval-only suite whose development partition is scored for every trial (out-of-domain check)")
440
+ ap.add_argument("--resume", action="store_true", help="finish evaluation for trials in --out that have a checkpoint but no result.json, and continue "
441
+ "unfinished full-weight trials from their last resume point (interrupted studies)")
442
+ a = ap.parse_args()
443
+ if a.aggregate:
444
+ aggregate(a.out); return
445
+ if a.resume:
446
+ if not a.suite: ap.error("--suite is required with --resume")
447
+ suite = Path(a.suite).resolve()
448
+ for directory in sorted(Path(a.out).iterdir()):
449
+ if (directory / "result.json").exists() or not (directory / "provenance.json").exists(): continue
450
+ transfer = Path(a.transfer).resolve() if a.transfer else None
451
+ if (directory / "checkpoint" / "head.pt").exists():
452
+ print(f"Resuming evaluation for {directory.name}", flush=True)
453
+ resume_trial(suite, directory, source_hashes(), a.device, transfer)
454
+ elif read_json(directory / "provenance.json")["config"].get("full_ft"):
455
+ print(f"Continuing training for {directory.name}", flush=True)
456
+ continue_trial(suite, directory, source_hashes(), a.device, transfer)
457
+ if (Path(a.out) / "results.jsonl").exists(): (Path(a.out) / "results.jsonl").unlink()
458
+ aggregate(a.out); return
459
+ if not a.plan or not a.suite:
460
+ ap.error("--suite and --plan are required unless --aggregate")
461
+ suite = Path(a.suite).resolve()
462
+ trials = load_plan(suite, a.plan)
463
+ if a.dry_run:
464
+ print(json.dumps({"trials": trials, "existing": a.existing, "suite_sha256": digest(suite / "manifest.json"), "locked_test": "not read"}, indent=2))
465
+ return
466
+ expected_sources = source_hashes()
467
+ with study_lock():
468
+ if a.wait_pid:
469
+ print(f"Waiting for existing training process {a.wait_pid}; no competing GPU job will start.", flush=True)
470
+ while True:
471
+ try:
472
+ os.kill(a.wait_pid, 0)
473
+ except ProcessLookupError:
474
+ break
475
+ time.sleep(10)
476
+ output = Path(a.out).resolve()
477
+ output.mkdir(parents=True, exist_ok=False)
478
+ entries = [(None, p) for p in a.existing] + [(t, None) for t in trials]
479
+ for i, (config, existing) in enumerate(entries):
480
+ label = Path(existing).name if existing else f"trial-{i}"
481
+ print(f"Starting {label}", flush=True)
482
+ try:
483
+ execute_trial(config or {}, suite, output / f"{i:02d}-{label}", expected_sources, a.device, existing, Path(a.transfer).resolve() if a.transfer else None)
484
+ except Exception as error:
485
+ with (output / "results.jsonl").open("a", encoding=ENCODING) as ledger:
486
+ ledger.write(json.dumps({"id": label, "status": "failed", "error": str(error)}) + "\n")
487
+ raise
488
+ aggregate(output)
489
+
490
+
491
+ if __name__ == "__main__":
492
+ main()
runtime/kev/full_ft.py ADDED
@@ -0,0 +1,454 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Full-weight training (`kev.train --full_ft 1`): the whole text backbone and the pointer head are trained, not a LoRA.
2
+
3
+ The backbone's working weights are bf16 (what the forward and backward run in, and what the checkpoint stores). Plain
4
+ bf16 AdamW updates at lr ~1e-5 round away: a bf16 weight near 0.02 moves in steps of ~1e-4. So MasterAdamW keeps an
5
+ fp32 master copy of every weight plus fp32 AdamW moments, updates the masters exactly (torch's fused AdamW) and writes
6
+ bf16(master) back after each step. Where those 12 bytes per parameter live is the whole memory plan:
7
+
8
+ - one GPU (`offload`): in host memory. A 27B backbone is ~51 GB of bf16 weights plus ~51 GB of bf16 gradients on the
9
+ H200, and ~310 GB of masters and moments in host RAM; each step streams one tensor at a time (gradient down, update on
10
+ the CPU, bf16 weights up), so no fp32 gradient copy of the whole model is ever held. Technique from AutoJev
11
+ (github.com/denis-pplx/autojev, src/autojev/optim.py, MIT: fp32 masters and AdamW in host memory beside bf16 weights on
12
+ one H200); the per-tensor streaming and the clipping folded into the host pass are ours.
13
+ - several GPUs (torchrun, FSDP2 `shard`): each rank keeps its shard's masters and moments next to the shard on its GPU
14
+ (a 27B over 8 H200s: ~6.4 GB of weights, ~6.4 GB of gradients and ~38 GB of optimizer state per GPU), so the step is a
15
+ GPU kernel and the data-parallel ranks train 8 micro-batches at once.
16
+
17
+ Gradients accumulate in bf16 over the micro-batches of a step, as AutoJev's did. The pointer head is small and fp32; it
18
+ is replicated on every rank and its gradient summed across ranks before the step.
19
+
20
+ What a run writes besides its final checkpoint: resume points (ResumeWriter: the fp32 optimizer state, ~307 GB for a
21
+ 27B, replaced as the run goes and removed at its end) and snapshots (SnapshotWriter: loadable bf16 checkpoints at
22
+ registered steps, the final checkpoint's files, ~51 GB for a 27B, kept).
23
+ """
24
+ import copy
25
+ import datetime
26
+ import math
27
+ import os
28
+ import shutil
29
+ import threading
30
+ import time
31
+ from pathlib import Path
32
+
33
+ import torch
34
+ import torch.distributed as dist
35
+ from torch.optim.adamw import adamw
36
+
37
+ from .device import sync
38
+ from .suite import read_json, write_json
39
+
40
+ MIN_TORCH = (2, 8) # FSDPModule.set_gradient_divide_factor (shard) and the FSDP2 behaviour measured in PR #122 / #125
41
+
42
+
43
+ def unsupported_torch(version=torch.__version__):
44
+ """Why this torch cannot train full weights, or None. pyproject allows torch >= 2.6; kev.full_ft needs MIN_TORCH."""
45
+ have = tuple(int(part) for part in version.split("+")[0].split(".")[:2])
46
+ if have < MIN_TORCH:
47
+ return f"--full_ft 1 needs torch >= {'.'.join(map(str, MIN_TORCH))} (FSDPModule.set_gradient_divide_factor); this is torch {version}"
48
+ return None
49
+
50
+
51
+ SHARD_SIZE = "5GB" # save_pretrained shards: model-00001-of-000NN.safetensors + model.safetensors.index.json
52
+
53
+
54
+ def local(t):
55
+ """The part of a (possibly FSDP2-sharded, DTensor) tensor this rank owns."""
56
+ return t.to_local() if hasattr(t, "to_local") else t
57
+
58
+
59
+ def sharded(t):
60
+ return hasattr(t, "to_local")
61
+
62
+
63
+ class MasterAdamW(torch.optim.Optimizer):
64
+ """AdamW on fp32 master copies of (bf16) parameters, in host memory (`offload`) or on the parameter's device. The
65
+ param groups hold the model's own parameters, so lr schedulers work unchanged. step() clips to `max_grad_norm` by
66
+ the global gradient norm (across ranks when sharded), then updates one tensor at a time and frees its gradient."""
67
+
68
+ def __init__(self, groups, lr, weight_decay, offload, max_grad_norm=1.0, betas=(0.9, 0.999), eps=1e-8):
69
+ super().__init__(groups, {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay})
70
+ self.max_grad_norm, self.grad_norm = max_grad_norm, None
71
+ params = [p for g in self.param_groups for p in g["params"]]
72
+ for p in params:
73
+ master = local(p).detach().to("cpu" if offload else local(p).device, torch.float32, copy=True)
74
+ self.state[p] = {"master": master, "exp_avg": torch.zeros_like(master), "exp_avg_sq": torch.zeros_like(master),
75
+ "step": torch.zeros((), device=master.device)}
76
+ self.pinned = offload and any(p.is_cuda for p in params)
77
+ if self.pinned: # two pinned staging slots per direction and dtype: tensor i+1 moves while the CPU updates tensor i
78
+ size = lambda dtype: max(local(p).numel() for p in params if p.dtype == dtype)
79
+ self.down = {d: [torch.empty(size(d), dtype=d, pin_memory=True) for _ in range(2)] for d in {p.dtype for p in params}}
80
+ self.up = {d: [torch.empty(size(d), dtype=d, pin_memory=True) for _ in range(2)] for d in self.down}
81
+ self.scratch = torch.empty(max(local(p).numel() for p in params), dtype=torch.float32)
82
+ self.uploads = [None, None]
83
+
84
+ def _clip_scale(self, params):
85
+ """Clip coefficient from the global L2 norm of all gradients (torch.nn.utils.clip_grad_norm_'s formula); sharded
86
+ gradients are summed over ranks, replicated ones counted once."""
87
+ sq = lambda ps: sum(torch.linalg.vector_norm(local(p.grad), dtype=torch.float32) ** 2 for p in ps)
88
+ total = sq([p for p in params if sharded(p)]) + 0.0
89
+ if dist.is_initialized() and torch.is_tensor(total): dist.all_reduce(total)
90
+ self.grad_norm = float((total + sq([p for p in params if not sharded(p)])) ** 0.5)
91
+ if not math.isfinite(self.grad_norm): # a NaN/inf gradient, even from a finite loss, stops here before any master moves (on every rank: the norm is global); min(1.0, nan) would otherwise apply it unclipped
92
+ raise RuntimeError(f"non-finite gradient norm {self.grad_norm}; refusing the optimizer step")
93
+ return min(1.0, self.max_grad_norm / (self.grad_norm + 1e-6)) if self.max_grad_norm else 1.0
94
+
95
+ def _fetch(self, p, slot):
96
+ """Start copying p's gradient to host memory; -> a callable returning it as fp32 once the copy has landed."""
97
+ g = local(p.grad)
98
+ if not self.pinned: return lambda: g.float() # masters beside the weights (or everything on the CPU)
99
+ buf = self.down[g.dtype][slot][:g.numel()].view_as(g)
100
+ buf.copy_(g, non_blocking=True); done = torch.cuda.Event(); done.record()
101
+ def ready():
102
+ done.synchronize()
103
+ return self.scratch[:g.numel()].view_as(g).copy_(buf)
104
+ return ready
105
+
106
+ def _put(self, p, master, slot):
107
+ """Write bf16(master) into the working weights (through a pinned slot, asynchronously, when offloaded)."""
108
+ w = local(p)
109
+ if not self.pinned: w.copy_(master); return
110
+ if self.uploads[slot] is not None: self.uploads[slot].synchronize() # the slot's previous upload has left
111
+ buf = self.up[w.dtype][slot][:w.numel()].view_as(w)
112
+ buf.copy_(master); w.copy_(buf, non_blocking=True)
113
+ self.uploads[slot] = torch.cuda.Event(); self.uploads[slot].record()
114
+
115
+ def state_dict(self):
116
+ """This rank's optimizer state for a resume point: the group hyperparameters (the lr scheduler moves them) and, per
117
+ parameter in order, its fp32 master, moments and step. The bf16 weights are not saved: they are bf16(master)."""
118
+ return {"groups": [{k: v for k, v in g.items() if k != "params"} for g in self.param_groups],
119
+ "state": [self.state[p] for g in self.param_groups for p in g["params"]]}
120
+
121
+ @torch.no_grad()
122
+ def load_state_dict(self, saved):
123
+ """Restore state_dict() in place (torch.optim.Optimizer's own would cast the fp32 masters to the parameters' bf16)
124
+ and write bf16(master) back into the weights, which reproduces them bit for bit."""
125
+ params = [p for g in self.param_groups for p in g["params"]]
126
+ if len(saved["state"]) != len(params) or len(saved["groups"]) != len(self.param_groups):
127
+ raise ValueError("resume point does not match this model's parameters")
128
+ for group, hyper in zip(self.param_groups, saved["groups"]): group.update(hyper)
129
+ for p, s in zip(params, saved["state"]):
130
+ for key, value in s.items():
131
+ if self.state[p][key].shape != value.shape: raise ValueError(f"resume point: {key} shape {tuple(value.shape)} for {tuple(self.state[p][key].shape)}")
132
+ self.state[p][key].copy_(value)
133
+ local(p).copy_(self.state[p]["master"])
134
+
135
+ @torch.no_grad()
136
+ def step(self):
137
+ work = [(g, p) for g in self.param_groups for p in g["params"] if p.grad is not None]
138
+ if dist.is_initialized():
139
+ for _, p in work:
140
+ if not sharded(p): dist.all_reduce(p.grad) # replicated (the pointer head): each rank holds its share
141
+ scale = self._clip_scale([p for _, p in work])
142
+ pending = self._fetch(work[0][1], 0) if work else None
143
+ for i, (group, p) in enumerate(work):
144
+ grad = pending()
145
+ if i + 1 < len(work): pending = self._fetch(work[i + 1][1], (i + 1) % 2)
146
+ if scale < 1: grad.mul_(scale)
147
+ s = self.state[p]
148
+ adamw([s["master"]], [grad], [s["exp_avg"]], [s["exp_avg_sq"]], [], [s["step"]], fused=True, amsgrad=False,
149
+ beta1=group["betas"][0], beta2=group["betas"][1], lr=group["lr"], weight_decay=group["weight_decay"],
150
+ eps=group["eps"], maximize=False)
151
+ self._put(p, s["master"], i % 2)
152
+ p.grad = None
153
+
154
+
155
+ # --- several GPUs -----------------------------------------------------------------------------------------------------
156
+
157
+ def init_distributed(device):
158
+ """-> (rank, world size). Under torchrun (WORLD_SIZE > 1) joins the process group (NCCL on CUDA, gloo on CPU for tests)
159
+ and makes this rank's GPU the current one, so "cuda" means it everywhere; otherwise (0, 1) and nothing is initialised.
160
+ Either way turns on the CUDA allocator's expandable segments before the first allocation: micro-batch lengths vary
161
+ from 200 to 8,000 tokens, and the first 27B probe on one H200 ran out of memory with 6.7 GB reserved but unused."""
162
+ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
163
+ world = int(os.environ.get("WORLD_SIZE", "1"))
164
+ if world == 1: return 0, 1
165
+ if device == "cuda": torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
166
+ # rank 0 alone writes the gathered 27B checkpoint (~51 GB) while the others wait at a barrier: NCCL's 10 min default is too short
167
+ dist.init_process_group("nccl" if device == "cuda" else "gloo", timeout=datetime.timedelta(minutes=60))
168
+ return dist.get_rank(), world
169
+
170
+
171
+ def shard(model):
172
+ """FSDP2 over the backbone: one unit per decoder layer plus the root (embeddings, final norm). Gradients are summed
173
+ over ranks, not averaged (the trainer already divides each micro-batch's loss by the step's global record count). The
174
+ pointer head stays replicated (MasterAdamW sums its gradient), starting from rank 0's initialisation."""
175
+ from torch.distributed.device_mesh import init_device_mesh
176
+ from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
177
+ mesh = init_device_mesh(str(model.device).split(":")[0], (dist.get_world_size(),))
178
+ policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
179
+ units = [*model.lm.layers, model.lm]
180
+ for unit in units: fully_shard(unit, mesh=mesh, mp_policy=policy)
181
+ for unit in units: unit.set_gradient_divide_factor(1.0)
182
+ for t in model.head.parameters(): dist.broadcast(t.data, 0)
183
+
184
+
185
+ def rank_share(items, rank, world):
186
+ """This rank's equal share of one epoch's shuffled items: padded (by wrapping around) to a multiple of the world size,
187
+ so every rank runs the same number of micro-batches and FSDP's collectives line up (as DistributedSampler does). The
188
+ padding repeats the first items of the shuffled order, whatever their length (--length_sort 1 does not come through
189
+ here: train.microbatch_plan pads each step from its shuffled records before cutting them)."""
190
+ if world == 1: return items
191
+ padded = items + items[: -len(items) % world]
192
+ return padded[rank::world]
193
+
194
+
195
+ def global_sum(values):
196
+ """Per-rank counters summed over the ranks (as they are on one process)."""
197
+ if not dist.is_initialized(): return values
198
+ total = torch.tensor(values, dtype=torch.float64, device="cuda" if dist.get_backend() == "nccl" else "cpu")
199
+ dist.all_reduce(total)
200
+ return total.tolist()
201
+
202
+
203
+ def global_max(values):
204
+ """Per-rank values, the largest over the ranks (every rank calls it: a collective)."""
205
+ if not dist.is_initialized(): return values
206
+ top = torch.tensor(values, dtype=torch.float64, device="cuda" if dist.get_backend() == "nccl" else "cpu")
207
+ dist.all_reduce(top, op=dist.ReduceOp.MAX)
208
+ return top.tolist()
209
+
210
+
211
+ # --- resume points ----------------------------------------------------------------------------------------------------
212
+
213
+ LATEST = "latest.json" # the resume point to continue from; written last (atomically) by rank 0, after every rank's file
214
+ WRITE_SHARE = 0.05 # at most this share of wall time may block on writing resume points
215
+ PEER_WAIT = 600 # seconds rank 0 waits, after writing its own file, for the other ranks' (they write in parallel)
216
+
217
+
218
+ def rank0_decides(value):
219
+ """`value` as rank 0 sees it, on every rank (a decision every rank must act on together); itself on one process."""
220
+ if not dist.is_initialized(): return value
221
+ flag = torch.tensor(float(value), device="cuda" if dist.get_backend() == "nccl" else "cpu")
222
+ dist.broadcast(flag, 0)
223
+ return bool(flag)
224
+
225
+
226
+ def save_due(step, every_steps, every_minutes, since, blocked):
227
+ """Whether to write a resume point after this optimizer step: every `every_steps` steps, or once `every_minutes`
228
+ have passed since `since`, stretched so the time training has blocked on writes (`blocked` seconds per point) stays
229
+ under WRITE_SHARE of the interval. Under torchrun rank 0's clock decides for every rank."""
230
+ interval = max(60 * every_minutes, blocked / WRITE_SHARE)
231
+ return rank0_decides(bool(every_steps and step % every_steps == 0) or bool(every_minutes and time.time() - since >= interval))
232
+
233
+
234
+ class ResumeWriter:
235
+ """Writes resume points: each rank's optimizer state (its shard's fp32 masters and moments), the lr scheduler and the
236
+ RNG states to `step-N/rank<r>.pt`, then rank 0's `latest.json` with the position (the data order, counters, the
237
+ arguments it must be resumed with), then removes earlier points. A 27B's state is ~307 GB and the runs volume writes it
238
+ at well under 1 GB/s, so where the state lives on the GPUs (FSDP2) `save` copies it to host memory and a thread
239
+ writes it while training goes on (the next `save`, and `wait`, join it first); where it already lives in host memory
240
+ (one GPU, offload) there is no room for a copy and `save` writes it before returning.
241
+ `after`: the SnapshotWriter; latest.json waits for a snapshot still being written (and is not written if it failed),
242
+ so a resume point is never committed before the snapshots of the steps it has passed: a continuation from it could
243
+ not write them again."""
244
+
245
+ def __init__(self, resume_dir, background):
246
+ self.dir, self.background, self.thread, self.host, self.error = Path(resume_dir), background, None, None, None
247
+ self.seconds = [] # how long each point took to write, in the background or not
248
+ self.rank, self.world = (dist.get_rank(), dist.get_world_size()) if dist.is_initialized() else (0, 1)
249
+
250
+ def save(self, step, opt, sched, position, after=None):
251
+ self.wait()
252
+ state = opt.state_dict()
253
+ if self.background: # host copies, allocated once and reused, so the optimizer may move on
254
+ tensors = [t for s in state["state"] for t in s.values()]
255
+ if self.host is None: self.host = [torch.empty(t.shape, dtype=t.dtype, device="cpu") for t in tensors]
256
+ for h, t in zip(self.host, tensors): h.copy_(t, non_blocking=True)
257
+ sync(tensors[0].device.type)
258
+ copies = iter(self.host)
259
+ state = {"groups": state["groups"], "state": [{k: next(copies) for k in s} for s in state["state"]]}
260
+ rng = {"torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state() if torch.cuda.is_initialized() else None}
261
+ payload = {"optimizer": state, "scheduler": copy.deepcopy(sched.state_dict()), "rng": rng}
262
+ if self.background:
263
+ self.thread = threading.Thread(target=self._write, args=(step, payload, position, after), daemon=True); self.thread.start()
264
+ else:
265
+ self._write(step, payload, position, after)
266
+
267
+ def wait(self):
268
+ """Join the background write; its error, if any, is raised here."""
269
+ if self.thread is not None: self.thread.join(); self.thread = None
270
+ if self.error is not None: raise RuntimeError("writing the resume point failed") from self.error
271
+
272
+ def _write(self, step, payload, position, after=None):
273
+ started = time.time()
274
+ try: self._write_point(step, payload, position, after); self.seconds.append(round(time.time() - started, 1))
275
+ except BaseException as error:
276
+ self.error = error
277
+ if not self.background: raise
278
+
279
+ def _write_point(self, step, payload, position, after=None):
280
+ target = self.dir / f"step-{step:07d}"
281
+ target.mkdir(parents=True, exist_ok=True)
282
+ torch.save(payload, target / f".rank{self.rank}.pt.tmp")
283
+ os.replace(target / f".rank{self.rank}.pt.tmp", target / f"rank{self.rank}.pt")
284
+ if self.rank: return
285
+ deadline = time.time() + PEER_WAIT
286
+ while missing := [r for r in range(self.world) if not (target / f"rank{r}.pt").exists()]: # the ranks share one filesystem
287
+ if time.time() > deadline:
288
+ raise TimeoutError(f"resume point {target.name}: rank(s) {missing} did not finish writing within {PEER_WAIT // 60} min of rank 0; "
289
+ "latest.json still names the previous point")
290
+ time.sleep(2)
291
+ if after is not None and (error := after.join()) is not None: # the training thread raises it too (SnapshotWriter.wait)
292
+ raise RuntimeError(f"a snapshot failed to write, so resume point {target.name} is not made the latest: a continuation from "
293
+ "the previous one writes that snapshot again") from error
294
+ write_json(self.dir / LATEST, {"dir": target.name, "step": step, **position}, atomic=True)
295
+ for old in self.dir.glob("step-*"): # only earlier points: a rank done with this one may have started the next
296
+ if int(old.name.removeprefix("step-")) < step: shutil.rmtree(old)
297
+
298
+
299
+ def load_resume(resume_dir, opt, sched, args):
300
+ """-> the saved position (or None when there is no resume point), after restoring this rank's optimizer state, the
301
+ weights, the scheduler and the RNG states. Refuses a point written with other arguments or another world size."""
302
+ resume_dir = Path(resume_dir)
303
+ if not (resume_dir / LATEST).exists(): return None
304
+ position = read_json(resume_dir / LATEST)
305
+ world, rank = dist.get_world_size() if dist.is_initialized() else 1, dist.get_rank() if dist.is_initialized() else 0
306
+ if position["world"] != world or position["args"] != args:
307
+ changed = sorted(k for k in set(args) | set(position["args"]) if args.get(k) != position["args"].get(k))
308
+ raise ValueError(f"resume point {resume_dir} was written by {position['world']} rank(s) with other arguments: {changed or 'world size'}")
309
+ saved = torch.load(resume_dir / position["dir"] / f"rank{rank}.pt", map_location="cpu", mmap=True, weights_only=True)
310
+ opt.load_state_dict(saved["optimizer"]); sched.load_state_dict(saved["scheduler"])
311
+ torch.set_rng_state(saved["rng"]["torch"])
312
+ if saved["rng"]["cuda"] is not None: torch.cuda.set_rng_state(saved["rng"]["cuda"])
313
+ return position
314
+
315
+
316
+ def gather_backbone(lm):
317
+ """Under FSDP2: the bf16 backbone's full state dict, gathered into rank 0's host memory (every rank must call this;
318
+ the others get {}). On one process None: save_pretrained reads the model itself."""
319
+ if not dist.is_initialized(): return None
320
+ from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
321
+ return get_model_state_dict(lm, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
322
+
323
+
324
+ def write_backbone(lm, out, state):
325
+ """save_pretrained of the bf16 backbone: config.json + model-*.safetensors (+ index); `state` from gather_backbone
326
+ (transformers strips the FSDP class prefix)."""
327
+ lm.save_pretrained(out, state_dict=state, max_shard_size=SHARD_SIZE)
328
+
329
+
330
+ def save_backbone(lm, out):
331
+ """The final checkpoint's backbone. Under FSDP2 every rank joins the gather and rank 0 writes while the others wait."""
332
+ state = gather_backbone(lm)
333
+ if not dist.is_initialized() or dist.get_rank() == 0: write_backbone(lm, out, state)
334
+ if dist.is_initialized(): dist.barrier()
335
+
336
+
337
+ # --- snapshots --------------------------------------------------------------------------------------------------------
338
+
339
+ SNAPSHOT_INFO = "snapshot.json" # written last (atomically) into a snapshot's checkpoint directory: the snapshot is complete
340
+
341
+
342
+ def snapshot_fractions(text):
343
+ """--snapshot_fractions "0.25,0.5,0.75" -> (0.25, 0.5, 0.75); "" or "none" -> (). Each strictly between 0 and 1 (the
344
+ final checkpoint is always written)."""
345
+ if str(text).strip().lower() in ("", "none"): return ()
346
+ try: values = tuple(sorted({float(part) for part in str(text).split(",")}))
347
+ except ValueError: raise ValueError(f"snapshot fractions are comma-separated numbers, or none: {text!r}") from None
348
+ if not all(0 < v < 1 for v in values): raise ValueError(f"snapshot fractions must lie strictly between 0 and 1: {text!r}")
349
+ return values
350
+
351
+
352
+ def snapshot_steps(total, fractions=(), every=0):
353
+ """The optimizer steps after which a run of `total` steps writes a snapshot: the first step at or past each fraction
354
+ of `total`, and every `every` steps; never the last step (that is the final checkpoint)."""
355
+ steps = {math.ceil(round(f * total, 6)) for f in fractions} | (set(range(every, total, every)) if every else set())
356
+ return sorted(s for s in steps if 0 < s < total)
357
+
358
+
359
+ def too_many_snapshots(fractions=(), every=0, total=None):
360
+ """Why this snapshot plan exceeds kev.budget.MAX_SNAPSHOTS (the disk a run's kept snapshots may take), or None.
361
+ `total`: the run's optimizer steps, or an upper bound on them (a trial's max_steps); None = unknown, which an
362
+ every-N plan cannot be checked against."""
363
+ from .budget import MAX_SNAPSHOTS
364
+ if every and total is None:
365
+ return "snapshot_every_steps needs the run's step count to be bounded (max_steps), so its snapshot count can be checked"
366
+ count = len(snapshot_steps(total, fractions, every)) if total is not None else len(fractions)
367
+ if count > MAX_SNAPSHOTS:
368
+ return (f"{count} snapshots planned (fractions {list(fractions)}, every {every} of {total} steps): at most kev.budget.MAX_SNAPSHOTS = "
369
+ f"{MAX_SNAPSHOTS} fit the container disk next to the resume points (a full disk fails the trial)")
370
+ return None
371
+
372
+
373
+ def snapshot_path(root, step):
374
+ return Path(root) / f"step-{step:07d}" / "checkpoint" # zero-padded like resume points (ResumeWriter)
375
+
376
+
377
+ def completed_snapshot_dirs(root):
378
+ """{step: checkpoint directory} of the complete snapshots under `root` (a snapshot is complete once its SNAPSHOT_INFO
379
+ exists); step-<N> directories are read by number, padded or not (the first snapshots, before padding, were not)."""
380
+ return {int(p.parent.parent.name.removeprefix("step-")): p.parent for p in Path(root).glob(f"step-*/checkpoint/{SNAPSHOT_INFO}")}
381
+
382
+
383
+ def completed_snapshots(root):
384
+ """Steps of the complete snapshots under `root`."""
385
+ return sorted(completed_snapshot_dirs(root))
386
+
387
+
388
+ class SnapshotWriter:
389
+ """Loadable bf16 checkpoints during a full-weight run, at registered optimizer steps, into `<root>/step-<N>/checkpoint`:
390
+ exactly the final checkpoint's files (save_pretrained shards + head.pt + tokenizer) plus SNAPSHOT_INFO, written last,
391
+ which marks the snapshot complete. Snapshots are kept (never deleted); a complete one is never written again (a
392
+ continued run skips it) and an incomplete one (a write the container did not finish) is replaced.
393
+ Under FSDP2 every rank joins the gather into rank 0's host memory (the only time training blocks), then rank 0 writes
394
+ from that copy in a thread while training goes on; the next `save`, `wait` and a resume point's latest.json join it
395
+ first. On one GPU the host holds the masters and there is no room for a copy, so `save` writes before returning."""
396
+
397
+ def __init__(self, root, steps, background):
398
+ self.root, self.steps, self.background, self.thread, self.error = Path(root), set(steps), background, None, None
399
+ self.rank = dist.get_rank() if dist.is_initialized() else 0
400
+ self.written = [] # this process's snapshots: {"step", "blocking_seconds", "write_seconds"}
401
+
402
+ def due(self, step):
403
+ """Whether to write a snapshot after this step: a registered step without a complete snapshot (rank 0 decides)."""
404
+ return step in self.steps and rank0_decides(step not in completed_snapshots(self.root))
405
+
406
+ def missed(self, step):
407
+ """Registered steps up to `step` (a resume point) without a complete snapshot: a continuation cannot write them."""
408
+ done = set(completed_snapshots(self.root))
409
+ return sorted(s for s in self.steps if s <= step and s not in done)
410
+
411
+ def save(self, step, lm, finish, info):
412
+ """Write the snapshot of this step: the backbone, then finish(directory) (head.pt and the tokenizer, as for the
413
+ final checkpoint), then SNAPSHOT_INFO with `info` and the time it took."""
414
+ self.wait()
415
+ started = time.time()
416
+ state = gather_backbone(lm)
417
+ if self.rank: return
418
+ blocked = time.time() - started
419
+ if self.background:
420
+ self.thread = threading.Thread(target=self._write, args=(step, lm, state, finish, info, blocked), daemon=True); self.thread.start()
421
+ else:
422
+ self._write(step, lm, state, finish, info, blocked)
423
+
424
+ def join(self):
425
+ """Wait for the snapshot being written, if any (any thread may); -> its error or None."""
426
+ thread = self.thread
427
+ if thread is not None: thread.join()
428
+ return self.error
429
+
430
+ def wait(self):
431
+ """Join the background write; its error, if any, is raised here, on every rank at once (every rank calls wait:
432
+ in save and at the end of training), so ranks 1..N do not go on into the next gather and hang on rank 0."""
433
+ self.join(); self.thread = None
434
+ if rank0_decides(self.error is not None):
435
+ raise RuntimeError("writing a snapshot failed" + ("" if self.error else " on rank 0")) from self.error
436
+
437
+ def _write(self, step, lm, state, finish, info, blocked):
438
+ started = time.time()
439
+ try:
440
+ target = snapshot_path(self.root, step)
441
+ if (target / SNAPSHOT_INFO).exists(): # complete (a racing writer, or an earlier attempt): never deleted or rewritten
442
+ print(f"snapshot step {step}: {target} is complete already; left as it is", flush=True); return
443
+ if target.exists(): shutil.rmtree(target) # an incomplete write
444
+ target.mkdir(parents=True)
445
+ write_backbone(lm, target, state)
446
+ finish(target)
447
+ timing = {"step": step, "blocking_seconds": round(blocked + (0 if self.background else time.time() - started), 1),
448
+ "write_seconds": round(time.time() - started, 1)}
449
+ write_json(target / SNAPSHOT_INFO, {**info, **timing}, atomic=True)
450
+ self.written.append(timing)
451
+ print(f"snapshot step {step}: {target} (training blocked {timing['blocking_seconds']} s, written in {timing['write_seconds']} s)", flush=True)
452
+ except BaseException as error:
453
+ self.error = error
454
+ if not self.background: raise
runtime/kev/fused_qwen35.py ADDED
@@ -0,0 +1,155 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Fused CUDA kernels for serving the Qwen3.5 decoder layers (flash-linear-attention's Triton ops).
2
+
3
+ With CUDA graphs and batching (kev.cuda_graphs) a served pass is limited by the GPU work itself, and transformers'
4
+ reference Qwen3.5 layers spend over a third of it outside the matrix multiplies: the causal convolution falls back to a
5
+ PyTorch depthwise conv behind a concatenation of the cached conv state, the gating, the query/key head repeat and the
6
+ gated norm run as a dozen separate fp32 elementwise kernels per DeltaNet layer, and every RMSNorm and SwiGLU is several
7
+ more. fuse() rewrites those layers in place for inference:
8
+ - DeltaNet mixer: one projection GEMM for q/k/v, z, b and a (their weights concatenated); fla's Triton causal conv
9
+ that starts from the cached conv state (no concatenation); fla's chunked gated delta rule with the gate
10
+ (-exp(A_log) * softplus(a + dt_bias)), the beta sigmoid and the q/k L2 norm computed inside the kernel and the key
11
+ heads shared by their value heads (no repeat); fla's fused gated RMSNorm.
12
+ - MLP: one GEMM for gate and up (weights concatenated) and a fused SwiGLU.
13
+ - the decoder layer's two zero-centred RMSNorms as fla's fused RMSNorm with weight 1 + w in fp32, the second one also
14
+ adding the residual.
15
+ - attention: one GEMM for q (with its output gate), k and v, the q/k RMSNorms fused, the sigmoid output gate fused.
16
+ The math is the reference's; the rounding is not (the fused kernels keep fp32
17
+ where the reference rounds to bf16 in between), so fused and reference answers agree to bf16 noise, like the graphs.
18
+ Only for serving: there is no backward, the fused projections replace the originals (the merged LoRA is inside), and a
19
+ pass that continues a cached DeltaNet state does not advance it (see deltanet_forward).
20
+ """
21
+ import types
22
+
23
+ import torch
24
+ import torch.nn.functional as F
25
+ from fla.modules.activations import sigmoidglu, swiglu
26
+ from fla.modules.conv import causal_conv1d
27
+ from fla.modules.fused_norm_gate import rms_norm_gated
28
+ from fla.modules.layernorm import rms_norm
29
+ from fla.ops.gated_delta_rule import chunk_gated_delta_rule
30
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
31
+ from transformers.models.qwen3_5.modeling_qwen3_5 import apply_mask_to_padding_states, apply_rotary_pos_emb, eager_attention_forward
32
+
33
+
34
+ class _FixedNB:
35
+ """A Triton kernel launched with NB=1. fla passes these kernels NB, a token-count class (cdiv of the rows by a few
36
+ thousand), as a constexpr and an autotuning key, but their bodies never read it: every new batch shape compiled and
37
+ autotuned them again, stalling a busy server for about a second each time."""
38
+
39
+ def __init__(self, kernel):
40
+ self.kernel = kernel
41
+
42
+ def __getitem__(self, grid):
43
+ launch = self.kernel[grid]
44
+ return lambda *args, **kwargs: launch(*args, **{**kwargs, "NB": 1})
45
+
46
+
47
+ def _fix_nb():
48
+ import fla.modules.conv.triton.ops as conv, fla.modules.fused_norm_gate as gated, fla.modules.l2norm as l2, fla.modules.layernorm as norm
49
+ for module, name in ((conv, "causal_conv1d_fwd_kernel"), (gated, "layer_norm_gated_fwd_kernel"), (l2, "l2norm_fwd_kernel"), (norm, "layer_norm_fwd_kernel")):
50
+ if not isinstance(getattr(module, name), _FixedNB): setattr(module, name, _FixedNB(getattr(module, name)))
51
+
52
+
53
+ FLA_VERSION = "0.5.2" # the flash-linear-attention these kernels and _fix_nb were written and measured against
54
+
55
+
56
+ def _concat(*linears):
57
+ """One weight [sum(out), in] for several bias-free projections of the same input."""
58
+ if any(l.bias is not None for l in linears): raise ValueError("fused projections assume bias-free Linear layers")
59
+ return torch.cat([l.weight for l in linears], 0).contiguous()
60
+
61
+
62
+ def deltanet_forward(self, hidden_states, cache_params=None, attention_mask=None, **kwargs):
63
+ """Qwen3_5GatedDeltaNet.forward with fused kernels. Cache contract: a pass on an empty cache (a state) fills it, as the
64
+ reference does; a pass continuing a cached state (question rows) reads it and leaves it as it was. Kev's question rows
65
+ never continue from each other, and writing back every row's final states cost as much as reading them."""
66
+ hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
67
+ B, T, _ = hidden_states.shape
68
+ mixed, z, b, a = F.linear(hidden_states, self.in_proj).split(self.splits, -1)
69
+ layer = cache_params.layers[self.layer_idx] if cache_params is not None else None
70
+ previous = layer is not None and layer.has_previous_state[0]
71
+ fill = layer is not None and not previous
72
+ mixed, conv_state = causal_conv1d(mixed, self.conv_weight, None, initial_state=layer.conv_states[0] if previous else None,
73
+ output_final_state=fill, activation="silu")
74
+ q, k, v = mixed.split([self.key_dim, self.key_dim, self.value_dim], -1)
75
+ out, recurrent = chunk_gated_delta_rule(
76
+ q.reshape(B, T, -1, self.head_k_dim), k.reshape(B, T, -1, self.head_k_dim), v.reshape(B, T, -1, self.head_v_dim),
77
+ g=a, beta=b, initial_state=layer.recurrent_states[0] if previous else None, output_final_state=fill,
78
+ use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, A_log=self.A_log, dt_bias=self.dt_bias, use_beta_sigmoid_in_kernel=True)
79
+ if fill:
80
+ if not layer.is_conv_states_initialized[0]: layer.lazy_initialization(conv_states=conv_state, conv_kernel_size=conv_state.shape[-1])
81
+ layer.conv_states[0].copy_(conv_state) # in place: graph buffers keep their address
82
+ layer.has_previous_state[0] = True
83
+ cache_params.update_recurrent_state(recurrent, self.layer_idx)
84
+ out = rms_norm_gated(out, z.reshape(B, T, -1, self.head_v_dim), self.norm.weight, None, activation="swish", eps=self.norm.variance_epsilon)
85
+ return self.out_proj(out.reshape(B, T, -1))
86
+
87
+
88
+ def attention_forward(self, hidden_states, position_embeddings, attention_mask, past_key_values=None, **kwargs):
89
+ """Qwen3_5Attention.forward with one projection GEMM, fused q/k RMSNorms and a fused sigmoid output gate."""
90
+ shape = hidden_states.shape[:-1]
91
+ q_gate, k, v = F.linear(hidden_states, self.qkv).split(self.splits, -1)
92
+ q, gate = q_gate.reshape(*shape, -1, 2 * self.head_dim).chunk(2, -1)
93
+ q = rms_norm(q, self.q_norm_weight, None, eps=self.q_norm.eps).transpose(1, 2)
94
+ k = rms_norm(k.reshape(*shape, -1, self.head_dim), self.k_norm_weight, None, eps=self.k_norm.eps).transpose(1, 2)
95
+ q, k = apply_rotary_pos_emb(q, k, *position_embeddings)
96
+ v = v.reshape(*shape, -1, self.head_dim).transpose(1, 2)
97
+ if past_key_values is not None: k, v = past_key_values.update(k, v, self.layer_idx)
98
+ attend = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, eager_attention_forward)
99
+ out, weights = attend(self, q, k, v, attention_mask, dropout=0.0, scaling=self.scaling, **kwargs) # [B, T, heads, dim]: transposed back already
100
+ return self.o_proj(sigmoidglu(gate.reshape(*shape, -1), out.reshape(*shape, -1))), weights
101
+
102
+
103
+ def mlp_forward(self, x):
104
+ gate, up = F.linear(x, self.gate_up).chunk(2, -1)
105
+ return self.down_proj(swiglu(gate, up))
106
+
107
+
108
+ def decoder_forward(self, hidden_states, position_embeddings, attention_mask=None, position_ids=None, past_key_values=None, **kwargs):
109
+ """Qwen3_5DecoderLayer.forward with fused norms: the post-mixer norm also adds the residual."""
110
+ residual = hidden_states
111
+ h = rms_norm(hidden_states, self.input_norm_weight, None, eps=self.input_layernorm.eps)
112
+ if self.block_type == "linear_attention":
113
+ h = self.linear_attn(hidden_states=h, cache_params=past_key_values, attention_mask=attention_mask, **kwargs)
114
+ else:
115
+ h, _ = self.self_attn(hidden_states=h, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values,
116
+ position_embeddings=position_embeddings, **kwargs)
117
+ h, residual = rms_norm(h, self.post_norm_weight, None, residual=residual, eps=self.post_attention_layernorm.eps, prenorm=True)
118
+ return residual + self.mlp(h)
119
+
120
+
121
+ @torch.no_grad()
122
+ def fuse(lm):
123
+ """Rewrite a Qwen3.5 text backbone (DecisionModel.lm, merged) in place for fused-kernel serving. Contract changes
124
+ against the reference: no backward, and a pass that continues a cached DeltaNet state leaves that state as it was
125
+ (Kev's question rows never continue from each other). Needs flash-linear-attention FLA_VERSION exactly: it patches
126
+ fla's kernel launches (_fix_nb)."""
127
+ import fla
128
+ if fla.__version__ != FLA_VERSION:
129
+ raise RuntimeError(f"kev's fused kernels need flash-linear-attention=={FLA_VERSION}, found {fla.__version__}; "
130
+ "install that version or load with LoadOptions(fused=False) (KEV_FUSED=0 for kev.serve)")
131
+ if not all(hasattr(layer.mlp, "gate_proj") for layer in lm.layers):
132
+ raise NotImplementedError("kev's fused kernels cover the dense Qwen3.5 MLP, not mixture-of-experts layers; load with fused=False (KEV_FUSED=0)")
133
+ _fix_nb()
134
+ for layer in lm.layers:
135
+ if layer.block_type == "linear_attention":
136
+ m = layer.linear_attn
137
+ m.in_proj = _concat(m.in_proj_qkv, m.in_proj_z, m.in_proj_b, m.in_proj_a)
138
+ m.splits = [m.conv_dim, m.value_dim, m.num_v_heads, m.num_v_heads]
139
+ m.conv_weight = m.conv1d.weight.squeeze(1).contiguous()
140
+ del m.in_proj_qkv, m.in_proj_z, m.in_proj_b, m.in_proj_a
141
+ m.forward = types.MethodType(deltanet_forward, m)
142
+ else:
143
+ m = layer.self_attn
144
+ m.qkv = _concat(m.q_proj, m.k_proj, m.v_proj)
145
+ m.splits = [m.q_proj.out_features, m.k_proj.out_features, m.v_proj.out_features]
146
+ m.q_norm_weight, m.k_norm_weight = 1.0 + m.q_norm.weight.float(), 1.0 + m.k_norm.weight.float()
147
+ del m.q_proj, m.k_proj, m.v_proj
148
+ m.forward = types.MethodType(attention_forward, m)
149
+ layer.mlp.gate_up = _concat(layer.mlp.gate_proj, layer.mlp.up_proj)
150
+ del layer.mlp.gate_proj, layer.mlp.up_proj
151
+ layer.mlp.forward = types.MethodType(mlp_forward, layer.mlp)
152
+ layer.input_norm_weight = 1.0 + layer.input_layernorm.weight.float()
153
+ layer.post_norm_weight = 1.0 + layer.post_attention_layernorm.weight.float()
154
+ layer.forward = types.MethodType(decoder_forward, layer)
155
+ return lm
runtime/kev/jev.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import json
3
+ import os
4
+ import re
5
+ import subprocess
6
+ from pathlib import Path
7
+
8
+ from kev.benchmark import evaluate_records
9
+ from kev.predictors import JevPredictor
10
+ from kev.suite import digest, load_split, read_manifest, write_json
11
+
12
+
13
+ def provision_key(scope):
14
+ result = subprocess.run(["vercel", "ai-gateway", "api-keys", "create", "--scope", scope,
15
+ "--name", "kev-jev-evaluation", "--limit", "1", "--refresh-period", "none",
16
+ "--expiration", "7d", "--no-color", "--non-interactive"], capture_output=True, text=True)
17
+ if result.returncode:
18
+ raise RuntimeError("Vercel key creation failed; inspect account permissions in the Vercel dashboard")
19
+ text = result.stdout + result.stderr
20
+ match = re.search(r"\bvck_[A-Za-z0-9_-]+", text)
21
+ if not match:
22
+ raise RuntimeError("Key was created but its output format was not recognized; no key output was logged")
23
+ return match.group()
24
+
25
+
26
+ def main():
27
+ ap = argparse.ArgumentParser()
28
+ ap.add_argument("--suite", required=True)
29
+ ap.add_argument("--out", required=True)
30
+ ap.add_argument("--provision-scope", help="Explicitly authorize a $1 non-renewing, seven-day key on this team")
31
+ ap.add_argument("--budget", type=float, default=0.1)
32
+ ap.add_argument("--max-calls", type=int, default=700)
33
+ ap.add_argument("--count-refusals", action="store_true",
34
+ help="count requests Jev refuses (HTTP 400/413/422, e.g. past its context) as rejected records (rejected.json, "
35
+ "report coverage) and continue; other errors still stop the read")
36
+ ap.add_argument("--attempts", type=int, default=4, help="tries per request on gateway 5xx errors (backoff 1, 2, 4, ... s, at most 30 s)")
37
+ a = ap.parse_args()
38
+ # a provisioned key carries a $1 limit of its own; a caller's key may run a larger panel (breadth-v1 dev: ~2,000 records)
39
+ if not 0 < a.budget <= 5 or not 1 <= a.max_calls <= 5000 or not 1 <= a.attempts <= 12:
40
+ ap.error("budget must be in (0, 5], max-calls in [1, 5000] and attempts in [1, 12]")
41
+ if Path(a.out).exists():
42
+ ap.error("output directory already exists")
43
+ records = load_split(a.suite, "development")
44
+ key = os.environ.get("AI_GATEWAY_API_KEY")
45
+ if not key and a.provision_scope:
46
+ key = provision_key(a.provision_scope)
47
+ if not key:
48
+ ap.error("Set AI_GATEWAY_API_KEY or explicitly select --provision-scope")
49
+ predictor = JevPredictor(key, a.budget, a.max_calls, count_refusals=a.count_refusals, attempts=a.attempts)
50
+ try:
51
+ heldout = read_manifest(a.suite)["holdout_sources"]
52
+ report, _ = evaluate_records(records, predictor, a.out, heldout_sources=tuple(heldout), skip_overlong=a.count_refusals)
53
+ report.update(suite_sha256=digest(Path(a.suite) / "manifest.json"), split="development", provider=predictor.accounting())
54
+ write_json(Path(a.out) / "report.json", report)
55
+ print(json.dumps({"clean": report["clean"], "variants": report["variants"], "provider": report["provider"]}, indent=2))
56
+ finally:
57
+ predictor.close()
58
+ if Path(a.out).exists():
59
+ write_json(Path(a.out) / "usage.json", predictor.accounting())
60
+
61
+
62
+ if __name__ == "__main__":
63
+ main()
runtime/kev/metrics.py ADDED
@@ -0,0 +1,426 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Scoring of benchmark rows: accuracy, calibration (ECE, Brier, NLL), selective prediction (coverage at an error budget,
2
+ risk-coverage curve, AURC, thresholds), temperature fitting and the record-clustered paired bootstrap.
3
+
4
+ Pure numpy over the row dicts that kev.benchmark.prediction_rows produces ({"p", "label", "type", "keys", "task", "source",
5
+ "variant", ..., optionally "logits"}). No torch, no model: usable on saved rows.json files.
6
+ """
7
+ import math
8
+ from collections import defaultdict
9
+
10
+ import numpy as np
11
+
12
+ EPSILON = 1e-9
13
+
14
+
15
+ def ece(conf, correct, bins=10):
16
+ conf, correct = np.asarray(conf), np.asarray(correct, dtype=float)
17
+ edges = np.linspace(0, 1, bins + 1); e = 0.0
18
+ for lo, hi in zip(edges[:-1], edges[1:]):
19
+ m = (conf >= lo) & (conf < hi) if hi < 1 else (conf >= lo) & (conf <= hi)
20
+ if m.any(): e += m.mean() * abs(correct[m].mean() - conf[m].mean())
21
+ return float(e)
22
+
23
+
24
+ def _check_temperature(temperature):
25
+ if not math.isfinite(temperature) or temperature <= 0:
26
+ raise ValueError("temperature must be finite and positive")
27
+
28
+
29
+ def _tempered_logits(row, temperature):
30
+ """(z - max z) / T from the recorded logits, or from the floored probabilities when the row has none."""
31
+ _check_temperature(temperature)
32
+ p = np.asarray(row["p"], dtype=float)
33
+ z = np.asarray(row["logits"], dtype=float) if "logits" in row else np.log(np.maximum(p, EPSILON))
34
+ if z.shape != p.shape or not np.isfinite(z).all():
35
+ raise ValueError("logits must be finite and match the option count")
36
+ return (z - z.max()) / temperature
37
+
38
+
39
+ def probabilities_at_temperature(row, temperature=1.0):
40
+ """The row's returned probabilities at T=1; otherwise the softmax of its tempered logits."""
41
+ _check_temperature(temperature)
42
+ if temperature == 1:
43
+ return np.asarray(row["p"], dtype=float)
44
+ p = np.exp(_tempered_logits(row, temperature))
45
+ return p / p.sum()
46
+
47
+
48
+ def nll_at_temperature(row, temperature=1.0):
49
+ """Exact from logits when recorded; otherwise from the floored probabilities."""
50
+ if "logits" in row:
51
+ z = _tempered_logits(row, temperature)
52
+ return float(np.log(np.exp(z).sum()) - z[row["label"]])
53
+ return -math.log(max(float(probabilities_at_temperature(row, temperature)[row["label"]]), EPSILON))
54
+
55
+
56
+ def _row_scores(row):
57
+ """The additive per-question quantities whose means metrics() reports (used by the paired bootstrap)."""
58
+ p = np.asarray(row["p"], dtype=float); y = row["label"]
59
+ conf, ok = float(p.max()), bool(p.argmax() == y)
60
+ return {"acc": float(ok), "nll": nll_at_temperature(row), "brier": float(((p - np.eye(len(p))[y]) ** 2).sum()), "mean_conf": conf,
61
+ "confident_error_rate": float(conf >= 0.9 and not ok), "coverage_at_0_9": float(conf >= 0.9), "confidence_bias": conf - float(ok)}
62
+
63
+
64
+ def metrics(rows, temperature=1.0):
65
+ if not rows:
66
+ raise ValueError("cannot score an empty population")
67
+ nll, acc, conf, brier, mae, rps = [], [], [], [], [], []
68
+ for row in rows:
69
+ p = probabilities_at_temperature(row, temperature)
70
+ y = row["label"]
71
+ target = np.eye(len(p))[y]
72
+ nll.append(nll_at_temperature(row, temperature))
73
+ acc.append(int(p.argmax() == y)); conf.append(float(p.max()))
74
+ brier.append(float(((p - target) ** 2).sum()))
75
+ if row["type"] == "score":
76
+ mae.append(abs(float(p @ np.arange(len(p))) - y))
77
+ rps.append(float(((p.cumsum()[:-1] - target.cumsum()[:-1]) ** 2).mean()))
78
+ result = {"n": len(rows), "nll": float(np.mean(nll)), "acc": float(np.mean(acc)),
79
+ "ece": ece(conf, acc), "brier": float(np.mean(brier)), "mean_conf": float(np.mean(conf))}
80
+ confidence, correct = np.asarray(conf), np.asarray(acc, dtype=bool)
81
+ high = confidence >= 0.9
82
+ result.update(confident_error_rate=float(np.mean(high & ~correct)), coverage_at_0_9=float(high.mean()),
83
+ accuracy_at_0_9=float(correct[high].mean()) if high.any() else None,
84
+ coverage_at_5pct_error=coverage_at_error(confidence, correct, 0.05), coverage_at_1pct_error=coverage_at_error(confidence, correct, 0.01),
85
+ aurc=area_under_risk_coverage(confidence, correct),
86
+ error_rate_at_0_9=float((~correct[high]).mean()) if high.any() else None,
87
+ # signed over-confidence (mean top probability minus accuracy) and errors within the top confidence bins;
88
+ # the sign is diagnostic: untrained readouts run positive, outcome-trained ones near zero or negative
89
+ confidence_bias=float(confidence.mean() - correct.mean()),
90
+ top_bins={str(t): {"n": int((confidence >= t).sum()), "errors": int(((confidence >= t) & ~correct).sum()),
91
+ "error_rate": float((~correct[confidence >= t]).mean()) if (confidence >= t).any() else None} for t in (0.9, 0.95, 0.99)})
92
+ result["selective"] = {}
93
+ for fraction in (0.5, 0.8):
94
+ cutoff = np.sort(confidence)[-max(1, math.ceil(len(rows) * fraction))]
95
+ selected = confidence >= cutoff
96
+ result["selective"][str(fraction)] = {"coverage": float(selected.mean()), "accuracy": float(correct[selected].mean()),
97
+ "confidence_cutoff": float(cutoff)}
98
+ if mae:
99
+ result.update(score_mae=float(np.mean(mae)), ranked_probability_score=float(np.mean(rps)))
100
+ return result
101
+
102
+
103
+ def coverage_at_error(confidence, correct, budget):
104
+ """Selective automation: the largest share of decisions that can be accepted, in descending confidence order, while
105
+ the empirical error among the accepted stays <= budget (jev-benchmarks' "coverage at a fixed error budget"). A model
106
+ whose probabilities are honest gets high coverage; one that is confidently wrong gets little, whatever its accuracy."""
107
+ if not math.isfinite(budget) or not 0 <= budget <= 1:
108
+ raise ValueError("error budget must be in [0, 1]")
109
+ _, accepted, errors = _risk_curve_arrays(confidence, correct)
110
+ ok = np.flatnonzero(errors <= budget * accepted)
111
+ return float(accepted[ok[-1]] / accepted[-1]) if len(ok) else 0.0
112
+
113
+
114
+ def _selective_inputs(confidence, correct):
115
+ confidence, correct = np.asarray(confidence, dtype=float), np.asarray(correct)
116
+ if confidence.ndim != 1 or correct.shape != confidence.shape:
117
+ raise ValueError("confidence and correctness must be equal-length vectors")
118
+ if not np.isfinite(confidence).all() or ((confidence < 0) | (confidence > 1)).any():
119
+ raise ValueError("confidence must be finite and in [0, 1]")
120
+ if not np.isin(correct, [False, True]).all():
121
+ raise ValueError("correctness must be boolean")
122
+ return confidence, correct.astype(bool)
123
+
124
+
125
+ def _risk_curve_arrays(confidence, correct):
126
+ confidence, correct = _selective_inputs(confidence, correct)
127
+ if not len(confidence):
128
+ return confidence, np.array([], dtype=int), np.array([], dtype=int)
129
+ order = np.argsort(-confidence, kind="stable")
130
+ confidence, errors = confidence[order], np.cumsum(~correct[order])
131
+ ends = np.r_[np.flatnonzero(confidence[1:] != confidence[:-1]), len(order) - 1]
132
+ return confidence[ends], ends + 1, errors[ends]
133
+
134
+
135
+ def risk_coverage_curve(confidence, correct):
136
+ thresholds, accepted, errors = _risk_curve_arrays(confidence, correct)
137
+ return [{"threshold": float(t), "accepted": int(n), "errors": int(e),
138
+ "coverage": float(n / accepted[-1]), "risk": float(e / n)}
139
+ for t, n, e in zip(thresholds, accepted, errors)]
140
+
141
+
142
+ def area_under_risk_coverage(confidence, correct):
143
+ _, accepted, errors = _risk_curve_arrays(confidence, correct)
144
+ if not len(accepted):
145
+ return 0.0
146
+ return float(np.sum(np.diff(np.r_[0, accepted]) * errors / accepted) / accepted[-1])
147
+
148
+
149
+ def select_threshold(confidence, correct, budget, min_accepted=1):
150
+ if not math.isfinite(budget) or not 0 <= budget <= 1 or min_accepted < 1:
151
+ raise ValueError("invalid error budget or minimum accepted count")
152
+ thresholds, accepted, errors = _risk_curve_arrays(confidence, correct)
153
+ ok = np.flatnonzero((errors <= budget * accepted) & (accepted >= min_accepted))
154
+ return float(thresholds[ok[-1]]) if len(ok) else None
155
+
156
+
157
+ def evaluate_threshold(confidence, correct, threshold):
158
+ confidence, correct = _selective_inputs(confidence, correct)
159
+ if threshold is not None and (not math.isfinite(threshold) or not 0 <= threshold <= 1):
160
+ raise ValueError("threshold must be in [0, 1] or None for abstain-all")
161
+ accepted = confidence >= threshold if threshold is not None else np.zeros(len(confidence), dtype=bool)
162
+ n, errors = int(accepted.sum()), int((accepted & ~correct).sum())
163
+ return {"threshold": threshold, "n": len(confidence), "accepted": n, "errors": errors,
164
+ "coverage": n / len(confidence) if len(confidence) else 0.0, "risk": errors / n if n else None}
165
+
166
+
167
+ def unknowable_report(rows):
168
+ """Confidence on records whose deciding evidence was removed (source 'unknowable') against their intact controls.
169
+ Accuracy on the unknowable records is meaningless by construction; what is scored is whether the model knows it
170
+ cannot know: mean max-probability and the share of records answered at >= 0.9."""
171
+ unk = [r for r in rows if r["source"] == "unknowable"]; ctl = [r for r in rows if r["source"] == "unknowable_control"]
172
+ if not unk: return None
173
+ conf = lambda rs: [float(max(r["p"])) for r in rs]
174
+ by_id = {r["id"]: r for r in ctl}
175
+ paired = [(max(r["p"]), max(by_id[r["control_id"]]["p"])) for r in unk if r.get("control_id") in by_id]
176
+ return {"n": len(unk), "mean_max_p": float(np.mean(conf(unk))), "share_at_0_9": float(np.mean([c >= 0.9 for c in conf(unk)])),
177
+ "control_mean_max_p": float(np.mean(conf(ctl))) if ctl else None, "control_share_at_0_9": float(np.mean([c >= 0.9 for c in conf(ctl)])) if ctl else None,
178
+ "control_acc": float(np.mean([int(np.argmax(r["p"]) == r["label"]) for r in ctl])) if ctl else None,
179
+ "paired_confidence_drop": float(np.mean([c - u for u, c in paired])) if paired else None,
180
+ "share_less_confident_than_control": float(np.mean([u < c for u, c in paired])) if paired else None}
181
+
182
+
183
+ def grouped_metrics(rows, key, temperature=1.0):
184
+ groups = defaultdict(list)
185
+ for row in rows:
186
+ groups[row[key]].append(row)
187
+ return {name: metrics(group, temperature) for name, group in sorted(groups.items())}
188
+
189
+
190
+ LENGTH_EDGES = (8192, 16384, 32768, 65536) # state-token edges of calibration_by_length (long-context calibration)
191
+ LENGTH_METRICS = ("acc", "ece", "brier", "confident_error_rate")
192
+
193
+
194
+ def length_buckets(edges=LENGTH_EDGES):
195
+ """[(name, lo, hi)]: state tokens in [lo, hi) (hi None = no upper bound). The disjoint buckets (under_8k, 8k_16k, 16k_32k,
196
+ 32k_64k, 64k_plus for the default edges), then the tails of every inner edge (8k_plus, 16k_plus, 32k_plus). Edges are
197
+ increasing multiples of 1,024, which name them."""
198
+ if not edges or list(edges) != sorted(set(edges)) or any(e <= 0 or e % 1024 for e in edges):
199
+ raise ValueError("length edges must be increasing positive multiples of 1024")
200
+ k = lambda n: f"{n // 1024}k"
201
+ bounds = [0, *edges, None]
202
+ disjoint = [(f"under_{k(hi)}" if lo == 0 else f"{k(lo)}_plus" if hi is None else f"{k(lo)}_{k(hi)}", lo, hi) for lo, hi in zip(bounds[:-1], bounds[1:])]
203
+ return disjoint + [(f"{k(lo)}_plus", lo, None) for lo in edges[:-1]]
204
+
205
+
206
+ def calibration_by_length(rows, lengths=None, edges=LENGTH_EDGES):
207
+ """{bucket: {"n", acc, ece, brier, confident_error_rate}} of rows (scored as given, at T=1: pass them served) split by
208
+ state-token count: the row's own `state_tokens` when it records one, else lengths[row id] ({record id: tokens}); a count
209
+ is the encoded state segment, the <state> token included (kev.model.encode; kev.rounds.state_lengths computes it). An
210
+ empty bucket has n 0 and None values. The one home of per-length calibration (kev.rounds panels with `by_length`)."""
211
+ def tokens(row):
212
+ if row.get("state_tokens") is not None: return row["state_tokens"]
213
+ if lengths is None or row["id"] not in lengths: raise KeyError(f"no state-token count for record {row['id']!r}")
214
+ return lengths[row["id"]]
215
+ counted = [(tokens(r), r) for r in rows]
216
+ out = {}
217
+ for name, lo, hi in length_buckets(edges):
218
+ group = [r for n, r in counted if n >= lo and (hi is None or n < hi)]
219
+ scored = metrics(group) if group else {}
220
+ out[name] = {"n": len(group), **{m: scored.get(m) for m in LENGTH_METRICS}}
221
+ return out
222
+
223
+
224
+ def tempered_row(row, temperature):
225
+ """The row as a calibrated predictor would have returned it: probabilities and logits at `temperature`, so metrics()
226
+ at T=1 scores it (per-row temperatures, e.g. out-of-fold, cannot go through metrics(rows, T)). The recorded
227
+ inference_temperature composes (served at T, tempered by T' -> T * T'), so raw_row always restores the T=1 logits."""
228
+ return {**row, "p": probabilities_at_temperature(row, temperature).tolist(), "logits": _tempered_logits(row, temperature).tolist(),
229
+ "inference_temperature": row.get("inference_temperature", 1.0) * temperature}
230
+
231
+
232
+ def raw_row(row):
233
+ """The row as the checkpoint would have returned it at T=1: undoes the recorded inference_temperature (served logits
234
+ are z / T, so z = logits * T), which is what fit_temperature requires. A row served raw is returned unchanged."""
235
+ temperature = row.get("inference_temperature", 1.0)
236
+ _check_temperature(temperature)
237
+ if temperature == 1.0:
238
+ return row
239
+ if "logits" not in row:
240
+ raise ValueError("cannot restore raw logits from a row that recorded none")
241
+ z = np.asarray(row["logits"], dtype=float) * temperature
242
+ p = np.exp(z - z.max())
243
+ return {**row, "p": (p / p.sum()).tolist(), "logits": z.tolist(), "inference_temperature": 1.0}
244
+
245
+
246
+ def recorded(row):
247
+ """The row with its inference temperature explicit: rows saved before benchmarks recorded one (None) were scored raw."""
248
+ return {**row, "inference_temperature": 1.0} if row.get("inference_temperature") is None else row
249
+
250
+
251
+ def served_at(rows, temperature):
252
+ """The scored rows as a predictor at `temperature` would have returned them, whatever temperature they were saved at
253
+ (raw logits are restored first, so a Hub checkpoint's served rows and a trial's raw rows are treated alike)."""
254
+ return [tempered_row(raw_row(recorded(r)), temperature) for r in scored_rows(rows)]
255
+
256
+
257
+ def served(fit_rows, eval_rows, **fit_kwargs):
258
+ """(temperature fitted on `fit_rows`' raw logits, `eval_rows` served at it): how a checkpoint is calibrated and read
259
+ everywhere a comparison is served-vs-served. fit_kwargs override TEMPERATURE_FIT key by key."""
260
+ temperature = fit_temperature(served_at(fit_rows, 1.0), **{**TEMPERATURE_FIT, **fit_kwargs})
261
+ return temperature, served_at(eval_rows, temperature)
262
+
263
+
264
+ def scored_rows(rows):
265
+ """The rows metrics are computed on: clean variants of knowable records."""
266
+ return [row for row in rows if row["variant"] == "clean" and row["source"] != "unknowable"]
267
+
268
+
269
+ # how every released temperature was fitted (scripts/calibrate_checkpoint.py) and how kev.calibrate fits a workload's:
270
+ # min mean NLL over a 121-point log grid on 0.25..4, every question weighted equally. Research trials (kev.experiment) and
271
+ # the kev-finetune skill keep the default 81-point grid: their temperatures are screening reports, not shipped values.
272
+ TEMPERATURE_FIT = {"aggregation": "micro", "points": 121}
273
+ TEMPERATURE_FIT_METHOD = f"min {TEMPERATURE_FIT['aggregation']} mean NLL over a {TEMPERATURE_FIT['points']}-point log grid 0.25..4"
274
+
275
+
276
+ def fit_temperature(rows, aggregation="macro", points=81):
277
+ """Temperature minimizing the (micro or per-task macro) mean NLL over a `points`-point log grid on 0.25..4.
278
+ Released checkpoints and workload fits use **TEMPERATURE_FIT."""
279
+ if aggregation not in ("micro", "macro"):
280
+ raise ValueError("invalid calibration aggregation")
281
+ clean = scored_rows(rows)
282
+ if not clean:
283
+ raise ValueError("cannot fit temperature without labelled calibration rows")
284
+ if any(row.get("inference_temperature", 1.0) != 1.0 for row in clean):
285
+ raise ValueError("fit temperature on raw logits, not previously calibrated outputs")
286
+ candidates = np.exp(np.linspace(np.log(0.25), np.log(4), points))
287
+ weights = np.ones(len(clean))
288
+ if aggregation == "macro":
289
+ counts = defaultdict(int)
290
+ for row in clean:
291
+ counts[row["task"]] += 1
292
+ weights = np.asarray([1.0 / counts[row["task"]] for row in clean])
293
+ losses = [np.average([nll_at_temperature(r, float(t)) for r in clean], weights=weights) for t in candidates]
294
+ return float(candidates[int(np.argmin(losses))])
295
+
296
+
297
+ def _source_groups(rows):
298
+ """dict[source, list[np.ndarray of row indices]]: one array per (source, group) cluster, in first-seen order."""
299
+ groups = defaultdict(list)
300
+ for i, row in enumerate(rows):
301
+ groups[(row["source"], row["group"])].append(i)
302
+ sources = defaultdict(list)
303
+ for (source, _), indices in groups.items():
304
+ sources[source].append(np.asarray(indices, dtype=int))
305
+ return sources
306
+
307
+
308
+ def cluster_resamples(rows, samples, seed):
309
+ """Index arrays of `samples` bootstrap resamples of `rows`: within each source, (source, group) clusters drawn with
310
+ replacement, so sibling questions and variants of one record move together. The resampling unit of every bootstrap
311
+ here (paired_bootstrap, cross_validated_temperature) and of the research scripts that bootstrap their own statistic."""
312
+ sources, rng = _source_groups(rows), np.random.default_rng(seed)
313
+ for _ in range(samples):
314
+ yield np.concatenate([grouped[i] for grouped in sources.values() for i in rng.integers(0, len(grouped), size=len(grouped))])
315
+
316
+
317
+ def grouped_folds(rows, folds, seed):
318
+ """Fold id per row, disjoint in (source, group) so sibling questions and variants of one record never straddle
319
+ train/test, assigned round-robin within each source so every fold sees every source."""
320
+ units = sorted({(row["source"], row["group"]) for row in rows})
321
+ if len(units) < folds:
322
+ raise ValueError("fewer distinct (source, group) units than folds")
323
+ rng = np.random.default_rng(seed)
324
+ fold_of, offset = {}, 0
325
+ for source in sorted({source for source, _ in units}):
326
+ own = [unit for unit in units if unit[0] == source]
327
+ for i, j in enumerate(rng.permutation(len(own))):
328
+ fold_of[own[j]] = (offset + i) % folds
329
+ offset += len(own)
330
+ return np.asarray([fold_of[(row["source"], row["group"])] for row in rows])
331
+
332
+
333
+ def out_of_fold_rows(rows, folds=5, seed=0, **fit_kwargs):
334
+ """Each fold's temperature is fit on the other folds (fit_temperature(**fit_kwargs), raw rows required) and applied
335
+ to its held-out rows, so the result scores like any calibrated prediction. Returns (scored rows in input order,
336
+ the per-fold temperatures)."""
337
+ if folds < 2:
338
+ raise ValueError("folds must be >= 2")
339
+ clean = scored_rows(rows)
340
+ fold_id = grouped_folds(clean, folds, seed)
341
+ temperatures, oof = [], list(clean)
342
+ for fold in range(folds):
343
+ held = fold_id == fold
344
+ temperature = fit_temperature([row for row, out in zip(clean, held) if not out], **fit_kwargs)
345
+ temperatures.append(temperature)
346
+ for i in np.flatnonzero(held):
347
+ oof[i] = tempered_row(clean[i], temperature)
348
+ return oof, temperatures
349
+
350
+
351
+ def cross_validated_temperature(rows, folds=5, seed=0, samples=1000, **fit_kwargs):
352
+ """Out-of-fold calibration report: out_of_fold_rows scored against the raw rows. The bootstrap resamples
353
+ source-stratified groups (the unit paired_bootstrap uses) and reports raw vs out-of-fold ECE with a 95% interval on
354
+ the paired delta; `separated` is True when that interval excludes zero."""
355
+ if samples < 1:
356
+ raise ValueError("samples must be >= 1")
357
+ clean = scored_rows(rows)
358
+ oof, temperatures = out_of_fold_rows(clean, folds, seed, **fit_kwargs)
359
+ keys = ("n", "ece", "brier", "nll", "confident_error_rate", "coverage_at_5pct_error")
360
+ raw, out_of_fold = metrics(clean), metrics(oof)
361
+ correct = np.asarray([np.argmax(row["p"]) == row["label"] for row in clean]) # argmax is temperature-invariant
362
+ conf = np.asarray([[max(row["p"]) for row in rows] for rows in (clean, oof)])
363
+ sources = _source_groups(clean)
364
+ values = np.asarray([[ece(conf[0][drawn], correct[drawn]), ece(conf[1][drawn], correct[drawn])] for drawn in cluster_resamples(clean, samples, seed)])
365
+ ci = {name: np.quantile(v, [0.025, 0.975]).tolist() for name, v in (("raw", values[:, 0]), ("out_of_fold", values[:, 1]), ("delta", values[:, 1] - values[:, 0]))}
366
+ return {"raw": {k: raw[k] for k in keys}, "out_of_fold": {k: out_of_fold[k] for k in keys}, "ece_ci95": ci,
367
+ "separated": bool(ci["delta"][1] < 0 or ci["delta"][0] > 0), "temperatures": temperatures,
368
+ "folds": folds, "seed": seed, "samples": samples, "groups": sum(len(v) for v in sources.values()),
369
+ "unit": "source-stratified (source, group); sibling questions and variants stay together"}
370
+
371
+
372
+ def paired_bootstrap(candidate, reference, samples=1000, seed=0, metric="nll", aggregation="macro"):
373
+ nonlinear = {"coverage_at_5pct_error", "coverage_at_1pct_error", "aurc", "ece"} # recomputed on every resample
374
+ additive = metric not in nonlinear # else a mean of _row_scores
375
+ if (additive and metric not in _row_scores({"p": [1.0], "label": 0})) or aggregation not in ("micro", "macro") or samples < 1:
376
+ raise ValueError("unsupported bootstrap metric, aggregation, or sample count")
377
+
378
+ def index(rows):
379
+ out = {}
380
+ for r in scored_rows(rows):
381
+ key = r["id"], r["question"]
382
+ if key in out:
383
+ raise ValueError("duplicate paired example")
384
+ out[key] = r
385
+ return out
386
+
387
+ a, b = index(candidate), index(reference)
388
+ if not a or a.keys() != b.keys():
389
+ raise ValueError("paired comparison requires identical complete clean examples")
390
+ keys = sorted(a)
391
+ for key in keys:
392
+ row, other = a[key], b[key]
393
+ if row["keys"] != other["keys"] or row["label"] != other["label"]:
394
+ raise ValueError("paired comparison labels or option order differ")
395
+ if any(row[field] != other[field] for field in ("source", "group", "task", "type")):
396
+ raise ValueError("paired comparison group or task metadata differ")
397
+ sources = _source_groups([a[key] for key in keys])
398
+ task_names = np.asarray([a[key]["task"] for key in keys])
399
+ statistics = []
400
+ for indexed in (a, b):
401
+ rows = [indexed[key] for key in keys]
402
+ conf = np.asarray([max(r["p"]) for r in rows])
403
+ correct = np.asarray([np.argmax(r["p"]) == r["label"] for r in rows])
404
+ values = np.asarray([_row_scores(r)[metric] for r in rows]) if additive else None
405
+ statistics.append((conf, correct, values))
406
+
407
+ def statistic(indices, data):
408
+ conf, correct, values = data
409
+ if values is not None:
410
+ return float(values[indices].mean())
411
+ conf, correct = conf[indices], correct[indices]
412
+ if metric == "aurc":
413
+ return area_under_risk_coverage(conf, correct)
414
+ if metric == "ece":
415
+ return ece(conf, correct)
416
+ return coverage_at_error(conf, correct, 0.05 if metric == "coverage_at_5pct_error" else 0.01)
417
+
418
+ def delta(indices):
419
+ parts = [indices] if aggregation == "micro" else [indices[task_names[indices] == t] for t in np.unique(task_names[indices])]
420
+ return float(np.mean([statistic(part, statistics[0]) - statistic(part, statistics[1]) for part in parts]))
421
+
422
+ values = [delta(drawn) for drawn in cluster_resamples([a[key] for key in keys], samples, seed)]
423
+ return {f"{aggregation}_{metric}_delta": delta(np.arange(len(keys))), "ci95": np.quantile(values, [0.025, 0.975]).tolist(),
424
+ "samples": samples, "groups": sum(len(v) for v in sources.values()), "aggregation": aggregation,
425
+ "unit": "source-stratified original record; sibling questions stay together",
426
+ "method": "paired cluster percentile bootstrap; full statistic recomputed in each resample"}
runtime/kev/mirror.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """A private Hugging Face mirror of full-weight checkpoints: long-term storage for the snapshots (kev.full_ft.SnapshotWriter)
2
+ and final checkpoints that full-weight trials keep on the Modal runs volume. The volume copy stays primary: reads and
3
+ benchmarks load /runs/... paths, and a failed upload only logs (loudly) and never fails training.
4
+
5
+ mirror("/runs/r19-27b/00-trial-0/snapshots/step-0000389/checkpoint", "jaredpalmer/kev-snapshots")
6
+ -> the directory's files at <study>/<trial>/step-0000389/ in the repo (a final checkpoint: <study>/<trial>/final/)
7
+
8
+ Refuses a repo that is not private (it creates a missing one private). The upload commit is recorded next to the files
9
+ on the volume: in a snapshot's snapshot.json under "hub", in a final checkpoint's hub.json. modal_app.run_mirror runs
10
+ this in its own CPU container with the HF token of the Modal secret it mounts (HF_TOKEN in the environment; never printed).
11
+ """
12
+ import datetime
13
+ import re
14
+ from pathlib import Path
15
+
16
+ from .full_ft import SNAPSHOT_INFO
17
+ from .suite import read_json, write_json
18
+
19
+ DEFAULT_REPO = "jaredpalmer/kev-snapshots"
20
+ RECORD = "hub.json" # a final checkpoint's upload record (a snapshot's goes into its snapshot.json)
21
+ IGNORED = ["resume/*", ".*", RECORD] # resume points are optimizer state; dotfiles are partial writes
22
+ ATTEMPTS = 2 # an upload is tried once more after a failure
23
+
24
+
25
+ def destination(path, root="/runs"):
26
+ """Where a checkpoint directory goes in the repo: <study>/<trial>/<step-N> for a snapshot, <...>/final otherwise."""
27
+ rel = Path(path).relative_to(root).as_posix()
28
+ if match := re.fullmatch(r"(.+)/snapshots/(step-\d+)/checkpoint", rel): return f"{match[1]}/{match[2]}"
29
+ return rel.removesuffix("/checkpoint") + "/final"
30
+
31
+
32
+ def is_snapshot(path):
33
+ return (Path(path) / SNAPSHOT_INFO).exists()
34
+
35
+
36
+ def recorded(path):
37
+ """The upload record a checkpoint directory carries, or None."""
38
+ path = Path(path)
39
+ if is_snapshot(path): return read_json(path / SNAPSHOT_INFO).get("hub")
40
+ return read_json(path / RECORD) if (path / RECORD).exists() else None
41
+
42
+
43
+ def record(path, entry):
44
+ path = Path(path)
45
+ if is_snapshot(path): write_json(path / SNAPSHOT_INFO, {**read_json(path / SNAPSHOT_INFO), "hub": entry}, atomic=True) # stays complete throughout
46
+ else: write_json(path / RECORD, entry, atomic=True)
47
+
48
+
49
+ def ensure_private(api, repo):
50
+ """Create `repo` private if it is missing; refuse (PermissionError) if it exists and is not private."""
51
+ api.create_repo(repo, repo_type="model", private=True, exist_ok=True)
52
+ if not api.repo_info(repo, repo_type="model").private:
53
+ raise PermissionError(f"{repo} is not a private repo; checkpoints are only mirrored to private repos")
54
+
55
+
56
+ def mirror(path, repo=DEFAULT_REPO, api=None, root="/runs", force=False, log=print):
57
+ """Upload one complete checkpoint directory (a snapshot with its snapshot.json, or a checkpoint with head.pt) to the
58
+ private `repo` and record the commit. -> the record, or None when it was not uploaded (every failure is logged loudly
59
+ and nothing is raised: the volume copy is primary). Skips a directory already recorded for this repo unless `force`."""
60
+ path = Path(path)
61
+ if not (path / "head.pt").exists() or (path.parent.parent.name == "snapshots" and not is_snapshot(path)):
62
+ log(f"!!! mirror: {path} is not a complete checkpoint (no head.pt, or a snapshot without {SNAPSHOT_INFO}); not uploaded"); return None
63
+ if not force and (previous := recorded(path)) and previous.get("repo") == repo:
64
+ log(f"mirror: {path} is at {repo}/{previous['path']} already (commit {previous['commit']})"); return previous
65
+ if api is None:
66
+ from huggingface_hub import HfApi
67
+ api = HfApi() # HF_TOKEN from the environment (the Modal secret)
68
+ dest = destination(path, root)
69
+ for attempt in range(1, ATTEMPTS + 1):
70
+ try:
71
+ ensure_private(api, repo)
72
+ info = api.upload_folder(repo_id=repo, repo_type="model", folder_path=str(path), path_in_repo=dest,
73
+ ignore_patterns=IGNORED, commit_message=f"mirror {dest}")
74
+ except PermissionError as error: # not a transient failure: never retried, never uploaded
75
+ log(f"!!! mirror of {path} refused: {error}"); return None
76
+ except Exception as error: # noqa: BLE001 - loud, retried once, never raised
77
+ log(f"!!! mirror of {path} to {repo} failed (attempt {attempt}/{ATTEMPTS}): {type(error).__name__}: {str(error)[:300]}")
78
+ continue
79
+ entry = {"repo": repo, "path": dest, "commit": info.oid, "uploaded_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds")}
80
+ record(path, entry)
81
+ log(f"mirror: {path} -> {repo}/{dest} (commit {info.oid})")
82
+ return entry
83
+ log(f"!!! mirror of {path} to {repo} gave up after {ATTEMPTS} attempts; the volume copy is unaffected (retry: modal_app.py::mirror_snapshots)")
84
+ return None
runtime/kev/mlx_model.py ADDED
@@ -0,0 +1,229 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Apple Silicon backend for the Qwen3.5 checkpoints: mlx-lm's Metal implementation of the hybrid backbone under Kev's
2
+ own encoder and pointer head.
3
+
4
+ MPS has no Gated DeltaNet kernels, so the PyTorch path runs reference code there (Kev-4B ~0.8 s per request). This module
5
+ runs the same computation on Metal through mlx-lm and keeps everything Kev-specific unchanged: `kev.model.encode` builds the
6
+ tokens, `rows_of` splits them into one causal row per question (the hybrid form the torch path uses too), and the readout is
7
+ the very same `PointerHead` (fp32, with the checkpoint's temperature) applied to the branch hidden states.
8
+
9
+ Same contract as `DecisionModel` for serving and scoring: encode / forward / probs / probs_and_prefix / probs_with_prefix /
10
+ head / dtype. Selected by `LoadOptions(backend="mlx")` (or "auto" on Apple Silicon) in `kev.checkpoint`; never by the
11
+ benchmark, whose reported numbers stay on the fp32 torch path. Parity against that path is measured in
12
+ tests/test_mlx.py (max |dp| and argmax flips on development records, prefix vs full pass, one question vs several).
13
+
14
+ Two ways in, both through kev.checkpoint: a LoRA checkpoint is its base's mlx-lm model (`load_base`) with the adapter
15
+ folded in (`merge_lora`); a full-weight checkpoint (kev.train --full_ft, scripts/merge_lora_checkpoint.py) is built from its
16
+ own config.json and shards (`load_full`), nothing merged, so the only resident copy is the weights themselves.
17
+ """
18
+ import json
19
+ from pathlib import Path
20
+
21
+ import mlx.core as mx
22
+ import numpy as np
23
+ import torch
24
+ import torch.nn.functional as F
25
+ from mlx.utils import tree_flatten
26
+ from mlx_lm.models import qwen3_5
27
+ from mlx_lm.models.cache import make_prompt_cache
28
+ from mlx_lm.utils import load_model
29
+
30
+ from .model import PointerHead, encode, probs_one, rows_of, rows_per_pass
31
+
32
+ CACHE_LIMIT = 1 << 30 # MLX's buffer cache keeps a buffer per new request shape; Kev-4B on an M5, 50 requests: 1.0 GB cached vs 3.7 GB unbounded, same latency
33
+ MLX_DTYPES = {"bfloat16": mx.bfloat16, "float32": mx.float32} # config.json's dtype names (kev.checkpoint.Checkpoint.saved_dtype) -> MLX
34
+ FULL_MODEL_TYPE = "qwen3_5_text" # what save_pretrained of transformers' Qwen3_5TextModel writes (Qwen3.5 and Qwen3.8 alike)
35
+ FULL_PREFIX = "language_model.model." # the text backbone inside mlx-lm's qwen3_5.Model; a full checkpoint's names carry no prefix
36
+ PREFILL_CHUNK = 1024 # state tokens per prefix pass (MLXDecisionModel.prefix); on an M5 1,024 / 2,048 / 4,096 / one pass take the same time, the peak grows with the chunk
37
+
38
+
39
+ def load_base(base_dir):
40
+ """mlx-lm's model of a Hub base snapshot, weights as stored (bf16 for the Qwen3.5 bases); the LoRA path merges into it."""
41
+ mx.set_cache_limit(CACHE_LIMIT)
42
+ return load_model(Path(base_dir))[0]
43
+
44
+
45
+ def load_full(ckpt_dir, shards, dtype):
46
+ """mlx-lm's Qwen3.5 model holding a full-weight checkpoint's backbone: the architecture from the checkpoint's own
47
+ config.json (the text model's config, model_type qwen3_5_text, which mlx-lm's qwen3_5.ModelArgs takes as its
48
+ text_config), every tensor from `shards` in `dtype` (config.json's name: "bfloat16" for every kev.train --full_ft run).
49
+ Nothing is merged or cast: the model holds exactly the saved values after mlx-lm's own layout rules for a
50
+ transformers export (qwen3_5 sanitize: the DeltaNet conv kernel moved to MLX's axis order and the zero-centred RMSNorm
51
+ weights stored as 1 + w, as for the Hub bases the LoRA path loads). Strict: every name and shape must match the model
52
+ exactly and every tensor must be `dtype`, else ValueError. The LM head is removed before any weight is materialized:
53
+ Kev reads hidden states, and the checkpoint carries none (vision tower and MTP were never loaded either)."""
54
+ ckpt_dir = Path(ckpt_dir)
55
+ config = json.loads((ckpt_dir / "config.json").read_text(encoding="utf-8"))
56
+ if config.get("model_type") != FULL_MODEL_TYPE or "text_config" in config:
57
+ raise ValueError(f"{ckpt_dir}: the MLX backend loads full-weight checkpoints of transformers' Qwen3_5TextModel "
58
+ f"(model_type {FULL_MODEL_TYPE!r}); config.json has model_type {config.get('model_type')!r}")
59
+ if dtype not in MLX_DTYPES: raise ValueError(f"{ckpt_dir}: no MLX dtype for weights stored as {dtype!r}")
60
+ want = MLX_DTYPES[dtype]
61
+ args = qwen3_5.ModelArgs.from_dict(config)
62
+ text = qwen3_5.TextModelArgs.from_dict(args.text_config)
63
+ pattern = ["full_attention" if (i + 1) % text.full_attention_interval == 0 else "linear_attention" for i in range(text.num_hidden_layers)]
64
+ if config.get("layer_types", pattern) != pattern: # mlx-lm places attention layers by full_attention_interval alone
65
+ raise ValueError(f"{ckpt_dir}: config.json's layer_types are not every {text.full_attention_interval}th layer full attention, "
66
+ f"the only layout mlx-lm's qwen3_5 builds")
67
+ mx.set_cache_limit(CACHE_LIMIT)
68
+ lm = qwen3_5.Model(args) # parameters are lazy: nothing is allocated until evaluated
69
+ if "lm_head" in lm.language_model: del lm.language_model["lm_head"] # untied bases (Qwen3.8-27B): 2.5 GB never read
70
+ weights = {}
71
+ for shard in shards:
72
+ for name, value in mx.load(str(shard)).items():
73
+ if FULL_PREFIX + name in weights: raise ValueError(f"{ckpt_dir}: tensor {name} is in more than one shard")
74
+ weights[FULL_PREFIX + name] = value
75
+ conv = [k for k in weights if k.endswith("linear_attn.conv1d.weight")]
76
+ if not conv or any(weights[k].ndim != 3 or weights[k].shape[1] != 1 for k in conv):
77
+ # sanitize infers "transformers layout, shift the norms" from the conv kernels; anything else would load wrong norms silently
78
+ raise ValueError(f"{ckpt_dir}: expected transformers' DeltaNet conv kernels [channels, 1, width]; this is not a kev.train --full_ft export")
79
+ weights = lm.language_model.sanitize(weights)
80
+ expected = {k: v.shape for k, v in tree_flatten(lm.parameters())}
81
+ missing, unexpected = sorted(set(expected) - set(weights)), sorted(set(weights) - set(expected))
82
+ wrong = sorted(k for k in set(expected) & set(weights) if tuple(weights[k].shape) != tuple(expected[k]))
83
+ cast = sorted(k for k, v in weights.items() if v.dtype != want)
84
+ if missing or unexpected or wrong or cast:
85
+ strip = lambda names: [n.removeprefix(FULL_PREFIX) for n in names[:2]]
86
+ raise ValueError(f"{ckpt_dir} does not match mlx-lm's Qwen3.5 text model: {len(missing)} tensors missing (e.g. {strip(missing)}), "
87
+ f"{len(unexpected)} unexpected (e.g. {strip(unexpected)}), {len(wrong)} with another shape (e.g. {strip(wrong)}), "
88
+ f"{len(cast)} not {dtype} (e.g. {strip(cast)})")
89
+ lm.load_weights(list(weights.items()), strict=True)
90
+ del weights
91
+ lm.eval()
92
+ mx.eval(lm.parameters())
93
+ mx.clear_cache()
94
+ return lm
95
+
96
+
97
+ def merge_lora(lm, adapter_dir, scale=1.0):
98
+ """Fold a PEFT adapter into the mlx-lm model's weights the way the torch path does: W + (B @ A) * alpha / r in fp32,
99
+ rounded once to the backbone dtype. `scale` is LoadOptions.lora_scale (WiSE-FT interpolation). Returns the tensor count."""
100
+ adapter_dir = Path(adapter_dir)
101
+ cfg = json.loads((adapter_dir / "adapter_config.json").read_text(encoding="utf-8"))
102
+ if cfg.get("trainable_token_indices"):
103
+ raise ValueError("the MLX backend does not carry trained token embeddings (special_embeddings checkpoints); use backend=torch")
104
+ alpha = cfg["lora_alpha"] / (cfg["r"] ** 0.5 if cfg.get("use_rslora") else cfg["r"])
105
+ weights = mx.load(str(adapter_dir / "adapter_model.safetensors"))
106
+ params = dict(tree_flatten(lm.parameters()))
107
+ merged = {}
108
+ with mx.stream(mx.cpu): # the GPU's fp32 matmul is a reduced-precision fast path (~1e-3 relative on an M5); the merge is one-time and must be exact
109
+ for name, a in weights.items():
110
+ if not name.endswith(".lora_A.weight"):
111
+ continue
112
+ stem = name[: -len(".lora_A.weight")]
113
+ # peft names the wrapped text model `base_model.model.<layers...>`; mlx-lm nests it as `language_model.model.<layers...>`
114
+ target = stem.replace("base_model.model.", "language_model.model.", 1) + ".weight"
115
+ if target not in params:
116
+ raise ValueError(f"adapter tensor {stem} has no weight in the mlx-lm model (looked for {target})")
117
+ base = params[target]
118
+ delta = (weights[stem + ".lora_B.weight"].astype(mx.float32) @ a.astype(mx.float32)) * (alpha * scale)
119
+ merged[target] = (base.astype(mx.float32) + delta).astype(base.dtype)
120
+ mx.eval(merged[target]) # one tensor at a time: one graph over every target peaks at ~2.9x the base weights and leaves ~2x of them in MLX's buffer cache (25.8 GB RSS for the 4B)
121
+ lm.load_weights(list(merged.items()), strict=False)
122
+ mx.eval(lm.parameters())
123
+ del params, weights # drop the pre-merge weights and the adapter before clearing the cache, or ~7 GB stay cached
124
+ mx.clear_cache() # hand the merge transients back to the OS; MLX keeps freed buffers otherwise
125
+ return len(merged)
126
+
127
+
128
+ class MLXDecisionModel:
129
+ """Prefill-only scorer: hidden states from mlx-lm, logits from the shared torch PointerHead."""
130
+ backend, device, hybrid, option_isolation = "mlx", "mlx", True, False
131
+ prefix_min_tokens = 0 # kev.serve caches the state prefix for every request: on Metal the branch-only pass is always the cheaper one
132
+
133
+ def __init__(self, lm, pad_id, head_dim=256):
134
+ self.lm = lm # mlx-lm qwen3_5.Model from load_base (+ merge_lora) or load_full
135
+ self.text = self.lm.language_model.model # Qwen3_5TextModel: embeddings -> layers -> final norm = `.model.last_hidden_state`
136
+ self.pad_id = pad_id
137
+ self.head = PointerHead(getattr(self.text.embed_tokens, "dims", self.text.embed_tokens.weight.shape[1]), dp=head_dim).eval()
138
+
139
+ @property
140
+ def dtype(self):
141
+ return str(getattr(self.text.embed_tokens, "activation_dtype", self.text.embed_tokens.weight.dtype)).removeprefix("mlx.core.")
142
+
143
+ def eval(self):
144
+ self.head.eval(); return self
145
+
146
+ def encode(self, tok, rec, **kw):
147
+ return encode(tok, rec, option_isolation=False, **kw)
148
+
149
+ def _hidden(self, rows, cache=None):
150
+ """[N, L, d] hidden states of right-padded token rows. Pads sit after every real token and both layer kinds are
151
+ causal (attention: causal mask; DeltaNet: a left-to-right recurrence), so no real token sees a pad."""
152
+ L = max(len(r) for r in rows)
153
+ ids = mx.array([r + [self.pad_id] * (L - len(r)) for r in rows], dtype=mx.int32)
154
+ h = self.text(ids, cache=cache)
155
+ mx.eval(h)
156
+ return h
157
+
158
+ def _logits(self, h, decide, opts):
159
+ """One question's logits through the fp32 pointer head (temperature included, eval mode)."""
160
+ idx = mx.array([decide, *opts], dtype=mx.int32)
161
+ picked = torch.from_numpy(np.asarray(h[idx].astype(mx.float32)))
162
+ with torch.no_grad():
163
+ return self.head(picked[0], picked[1:])
164
+
165
+ def forward_rows(self, enc):
166
+ """Row form, as the torch path computes it: every question is one causal row of state + branch tokens, the state
167
+ recomputed per row. The reference the prefix form is checked against (tests/test_mlx.py); serving uses `forward`."""
168
+ S, _, rows = rows_of(enc)
169
+ chunk, out = rows_per_pass([S + r["ids"] for r in rows]), []
170
+ for start in range(0, len(rows), chunk):
171
+ part = rows[start:start + chunk]
172
+ h = self._hidden([S + r["ids"] for r in part])
173
+ out += [self._logits(h[i], len(S) + r["decide"], [len(S) + o for o in r["opts"]]) for i, r in enumerate(part)]
174
+ return out
175
+
176
+ # --- state prefix: the state runs once into an mlx-lm prompt cache (KV for the attention layers, conv + recurrent state
177
+ # for the DeltaNet layers); the branches run as one batch on a replicated copy, so the prefix stays pristine and can be
178
+ # reused by the next request with the same state. On Metal this is also the cheapest way to answer a single request
179
+ # (state once instead of once per question), so it is the only path `forward` / `probs` take.
180
+
181
+ def prefix(self, enc):
182
+ """The state into a fresh prompt cache, PREFILL_CHUNK tokens per pass, each pass's cache evaluated before the next
183
+ (mlx-lm's own prefill does the same; the hidden states a pass returns are never read). One pass over the whole
184
+ state kept every layer's activations of every token alive at once and left the DeltaNet conv states lazy (each
185
+ holding its layer's whole [Ls, conv_dim] input): Kev-4B, 8,192 tokens, 4.6 GB above the weights instead of 1.2 GB,
186
+ the same 6.5 s. Exact: a pass continues the cached keys/values (rotary offset = cache offset), conv window and
187
+ fp32 recurrent state, so chunked and one-pass differ by float reassociation only (fp32 on the CPU: max |dp| 1e-6;
188
+ bf16 on Metal: up to ~0.01, like any change of kernel shapes; tests/test_mlx.py, runs/mlx-long-states/prefill-ab-*)."""
189
+ Ls = enc["seg"].count(0)
190
+ cache, ids = make_prompt_cache(self.lm), enc["ids"][:Ls]
191
+ for start in range(0, Ls, PREFILL_CHUNK):
192
+ self.text(mx.array([ids[start:start + PREFILL_CHUNK]], dtype=mx.int32), cache=cache)
193
+ mx.eval([c.state for c in cache])
194
+ return Ls, cache
195
+
196
+ def _branch_logits(self, enc, cache):
197
+ """Branches as rows on a replicated copy of the state cache, rows_per_pass rows (and cache copies) at a time."""
198
+ _, _, rows = rows_of(enc)
199
+ chunk, out = rows_per_pass([r["ids"] for r in rows], enc["seg"].count(0)), []
200
+ for start in range(0, len(rows), chunk):
201
+ part = rows[start:start + chunk]
202
+ batch = [type(c).merge([c] * len(part)) for c in cache] # merge copies the arrays: `cache` is not mutated
203
+ h = self._hidden([r["ids"] for r in part], batch)
204
+ out += [self._logits(h[i], r["decide"], r["opts"]) for i, r in enumerate(part)]
205
+ return out
206
+
207
+ def _branch_probs(self, enc, cache):
208
+ return [F.softmax(z, -1) for z in self._branch_logits(enc, cache)]
209
+
210
+ def forward(self, enc):
211
+ """List of logits tensors, one per question."""
212
+ return self._branch_logits(enc, self.prefix(enc)[1])
213
+
214
+ def probs(self, enc):
215
+ return self._branch_probs(enc, self.prefix(enc)[1])
216
+
217
+ def probs_and_prefix(self, enc):
218
+ prefix = self.prefix(enc)
219
+ return self._branch_probs(enc, prefix[1]), prefix
220
+
221
+ def probs_with_prefix(self, enc, prefix):
222
+ Ls, cache = prefix
223
+ if enc["seg"].count(0) != Ls: raise ValueError("prefix does not match this record's state")
224
+ return self._branch_probs(enc, cache)
225
+
226
+ def probs_batch(self, encs, prefixes, keep):
227
+ """kev.serve's batch call: one request at a time on Metal."""
228
+ out = [probs_one(self, e, p, k) for e, p, k in zip(encs, prefixes, keep)]
229
+ return [o[0] for o in out], [o[1] for o in out]
runtime/kev/model.py ADDED
@@ -0,0 +1,508 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Decision model: causal LM backbone + block-causal branch mask + pointer readout."""
2
+ import copy, math, os, re
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+ from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer, DynamicCache
7
+ from transformers.cache_utils import LinearAttentionCacheLayerMixin
8
+
9
+ # Reuse existing rarely-used Qwen special tokens as delimiters (state, q, opt, /opt, decide) so no
10
+ # embedding rows need to be added/trained; LoRA adapts their meaning.
11
+ SPECIAL = ["<|fim_prefix|>", "<|fim_middle|>", "<|box_start|>", "<|box_end|>", "<|fim_suffix|>"]
12
+ # training context: state tokens, tokens per question branch, and the whole packed record. Frozen suites are admitted with
13
+ # this rule (kev.suite) and training applies it to records built on the fly, so train and eval see the same population.
14
+ MAX_STATE, MAX_BRANCH, MAX_PACKED = 384, 1024, 2048
15
+ # serving context (kev.serve): a state of up to 64k tokens (twice Jev's 32k; the Qwen3.5 / Qwen3.8 bases' window is 262k)
16
+ # and a question row (state + its branch, the encoder's max_branch) of up to 8k tokens more. Until 64k states these were
17
+ # 8,192 / 8,192 (SERVE_MAX_*_8K below).
18
+ SERVE_MAX_STATE = 65536
19
+ SERVE_MAX_BRANCH = SERVE_MAX_STATE + 8192
20
+ SERVE_MAX_PACKED = SERVE_MAX_STATE + SERVE_MAX_BRANCH
21
+ # tokens one inference pass holds (rows_per_pass), and the longest request an attention-only backbone runs packed
22
+ # (rows_form: its block-causal mask is L x L); longer ones run as rows
23
+ ROW_PASS_TOKENS = 16384
24
+ # the longest state a checkpoint may be trained on (kev.train --max_state): the longest one served, which still leaves
25
+ # every question its training branch budget (training_context(MAX_TRAIN_STATE)["max_branch"] <= SERVE_MAX_BRANCH)
26
+ MAX_TRAIN_STATE = SERVE_MAX_STATE
27
+ # the limits before 64k states: the suites frozen until then were admitted under these and record them (kev.suite.SERVING_CONTEXT_8K)
28
+ SERVE_MAX_STATE_8K, SERVE_MAX_BRANCH_8K, MAX_TRAIN_STATE_8K = 8192, 8192, 7552
29
+
30
+
31
+ def training_context(max_state=MAX_STATE):
32
+ """The encoder limits for training with the state limit lifted to `max_state`: the row (state + one branch) and
33
+ packed limits grow by the same amount, so every question keeps its token budget. training_context() is the default
34
+ training context (kev.suite.CONTEXT without `truncate`)."""
35
+ if not MAX_STATE <= max_state <= MAX_TRAIN_STATE:
36
+ raise ValueError(f"max_state must be in [{MAX_STATE}, {MAX_TRAIN_STATE}]")
37
+ extra = max_state - MAX_STATE
38
+ return {"max_state": max_state, "max_branch": MAX_BRANCH + extra, "max_packed": MAX_PACKED + extra}
39
+
40
+
41
+ def rows_per_pass(rows, prefix_len=0, budget=ROW_PASS_TOKENS):
42
+ """How many causal rows one inference forward pass takes: as many as fit `budget` tokens counting the cached state
43
+ each row carries (prefix_len) plus its own tokens, at least one. Memory per pass is bounded by one maximal row however
44
+ many questions a request has, and the rows are independent, so the answers do not depend on the split. Kev-4B on MLX,
45
+ a 4.8k-token state with 64 questions: 24.7 GB peak in one pass, 8.9 GB one row at a time (and faster: the cost was
46
+ 64 copies of the state cache)."""
47
+ return max(1, budget // (prefix_len + max(len(r) for r in rows)))
48
+
49
+
50
+ class ContextOverflow(ValueError):
51
+ """A record does not encode within its context (state, branch or packed limit). Serving turns it into a 422; the
52
+ benchmark counts it as a rejected record for suites scored as published (skip_overlong). When the state is what
53
+ overflows, state_tokens (its count, the <state> token included) and max_state (the limit) say by how much."""
54
+
55
+ def __init__(self, message, state_tokens=None, max_state=None):
56
+ super().__init__(message)
57
+ self.state_tokens, self.max_state = state_tokens, max_state
58
+
59
+
60
+ def load_tokenizer(name, revision=None):
61
+ return AutoTokenizer.from_pretrained(name, revision=revision)
62
+
63
+
64
+ def pad_id(tok):
65
+ """The id used to right-pad token rows (never attended to); Qwen tokenizers define one, others fall back to 0."""
66
+ return tok.pad_token_id if tok.pad_token_id is not None else 0
67
+
68
+
69
+ def is_hybrid(config):
70
+ """Whether a (text) config has Gated DeltaNet layers (Qwen3.5). Such backbones cannot honour the block-causal mask and
71
+ run the row form; on Apple Silicon they are what the MLX backend is for."""
72
+ return "linear_attention" in set(getattr(config, "layer_types", None) or [])
73
+
74
+
75
+ _SPECIAL_RE = re.compile(r"<\|([A-Za-z0-9_]+)\|>")
76
+
77
+
78
+ def user_tokens(tok, text):
79
+ """Tokenize caller-supplied text so it can never produce delimiter/control tokens (option boundaries are unforgeable).
80
+ The fast tokenizer ignores split_special_tokens, so `<|name|>` is rewritten to `<¦name¦>` before tokenizing."""
81
+ return tok(_SPECIAL_RE.sub(r"<¦\1¦>", text), add_special_tokens=False).input_ids
82
+
83
+
84
+ OPT_NONE, OPT_DECIDE = -1, -2 # values of enc["opt"]: instruction/state tokens, and the <decide> token
85
+
86
+
87
+ def encode(tok, rec, max_state=MAX_STATE, max_branch=MAX_BRANCH, strict=False, option_isolation=False):
88
+ """Pack one record: [<state> ...] then per-question [<q> instr <opt> o </opt>... <decide>].
89
+
90
+ Returns ids, seg (0 = state, k = question k), pos (branch positions restart after state),
91
+ decide_idx [Q], opt_idx [Q][K] (index of </opt> token for each option), opt (per-token option index within its
92
+ question: OPT_NONE for state/instruction, 0..K-1 for option spans, OPT_DECIDE for <decide>).
93
+
94
+ option_isolation=True: every option span is its own sub-branch (it sees state + instruction + itself only), all
95
+ option spans share the same position ids, and <decide> sits at one fixed position after the longest span. Then the
96
+ per-option representations and <decide>'s attention over them are permutation-invariant by construction.
97
+
98
+ Not strict, a state over max_state tokens is cut to its first max_state (state_truncated; state_tokens is the whole
99
+ state's count, the <state> token included, as the limit counts it). Strict, it raises ContextOverflow.
100
+ """
101
+ state_tokens = user_tokens(tok, rec["state"])
102
+ if strict and len(state_tokens) + 1 > max_state:
103
+ raise ContextOverflow(f"state exceeds {max_state} tokens: {len(state_tokens) + 1}", state_tokens=len(state_tokens) + 1, max_state=max_state)
104
+ S = [tok.convert_tokens_to_ids(SPECIAL[0])] + state_tokens[: max_state - 1]
105
+ ids, seg, pos, opt = list(S), [0] * len(S), list(range(len(S))), [OPT_NONE] * len(S)
106
+ q_id, o_id, c_id, d_id = (tok.convert_tokens_to_ids(t) for t in SPECIAL[1:])
107
+ decide_idx, opt_idx = [], []
108
+ for k, q in enumerate(rec["questions"], start=1):
109
+ instr = [q_id] + user_tokens(tok, q["instr"])
110
+ spans = [[o_id] + user_tokens(tok, o) + [c_id] for o in q["options"]]
111
+ br = instr + [t for sp in spans for t in sp] + [d_id]
112
+ if len(br) > max_branch - len(S):
113
+ raise ContextOverflow(f"branch too long: {len(br)} tokens with a {len(S)}-token state (row limit {max_branch})")
114
+ base = len(ids); p0 = len(S)
115
+ br_opt = [OPT_NONE] * len(instr) + [j for j, sp in enumerate(spans) for _ in sp] + [OPT_DECIDE]
116
+ if option_isolation:
117
+ longest = max(len(sp) for sp in spans)
118
+ br_pos = list(range(p0, p0 + len(instr))) + [p0 + len(instr) + i for sp in spans for i in range(len(sp))] + [p0 + len(instr) + longest]
119
+ else:
120
+ br_pos = list(range(p0, p0 + len(br)))
121
+ ends, cursor = [], len(instr)
122
+ for sp in spans:
123
+ cursor += len(sp); ends.append(cursor - 1)
124
+ ids += br; seg += [k] * len(br); pos += br_pos; opt += br_opt
125
+ decide_idx.append(base + len(br) - 1); opt_idx.append([base + e for e in ends])
126
+ return {"ids": ids, "seg": seg, "pos": pos, "opt": opt, "option_isolation": option_isolation, "decide_idx": decide_idx, "opt_idx": opt_idx,
127
+ "labels": [q["label"] for q in rec["questions"]], "state_tokens": len(state_tokens) + 1, "state_truncated": len(state_tokens) + 1 > max_state}
128
+
129
+
130
+ def admit(model, tok, rec, truncate=False):
131
+ """The serving admission check (kev.serve and the Space; the torch and MLX models both encode through encode above):
132
+ -> rec encoded within the serving context. A state over SERVE_MAX_STATE tokens (the <state> token included) raises
133
+ ContextOverflow saying how long it is and how to fix it, unless truncate=True: then its first SERVE_MAX_STATE tokens
134
+ are read and the encoding says so (state_truncated, state_tokens). A question whose row (state + its branch) is over
135
+ SERVE_MAX_BRANCH raises ContextOverflow either way. Benchmarks do not truncate either: kev.predictors.LocalPredictor
136
+ encodes strictly within its suite's context."""
137
+ try:
138
+ return model.encode(tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=not truncate)
139
+ except ContextOverflow as e:
140
+ if e.max_state is None: raise
141
+ raise ContextOverflow(f"state is {e.state_tokens:,} tokens, over the {e.max_state:,}-token limit (the <state> token included): "
142
+ "shorten the document or split it across requests", state_tokens=e.state_tokens, max_state=e.max_state) from None
143
+
144
+
145
+ def fits(rec, *tokenizers, max_state=MAX_STATE, max_branch=MAX_BRANCH, max_packed=MAX_PACKED):
146
+ """True when the internal record encodes strictly (no truncation) within the training context under every tokenizer
147
+ given (frozen suites are admitted against the tokenizers of all their pinned bases)."""
148
+ try:
149
+ return all(len(encode(tok, rec, max_state=max_state, max_branch=max_branch, strict=True)["ids"]) <= max_packed for tok in tokenizers)
150
+ except ValueError:
151
+ return False
152
+
153
+
154
+ def branch_mask(seg, device, dtype=torch.float32):
155
+ """attend(i,j) iff j<=i and (seg[j]==0 or seg[j]==seg[i]). Returns additive [1,1,L,L]."""
156
+ return branch_mask_batch([seg], device, dtype)
157
+
158
+
159
+ def branch_mask_batch(segs, device, dtype=torch.float32, opts=None, length=None):
160
+ """Batched block-causal mask, additive [B,1,L,L], right-padded to the longest sequence.
161
+
162
+ Padded key positions are masked for every query; padded query rows keep the diagonal so no row is fully
163
+ masked (finfo.min, not -inf, so softmax stays finite either way). Real tokens never see pads because pads sit
164
+ after them (causal) and belong to no segment (-1).
165
+
166
+ opts (option isolation): within a question, an option-span token may attend to state, the instruction, and its own
167
+ span only; <decide> attends to everything in its question. Instruction tokens never see option spans (causal)."""
168
+ L = max(max(len(s) for s in segs), length or 0)
169
+ s = torch.full((len(segs), L), -1, device=device)
170
+ for b, seg in enumerate(segs):
171
+ s[b, : len(seg)] = torch.tensor(seg, device=device)
172
+ causal = torch.tril(torch.ones(L, L, dtype=torch.bool, device=device))
173
+ same = (s[:, None, :] == s[:, :, None]) | (s[:, None, :] == 0)
174
+ valid_key = (s != -1)[:, None, :]
175
+ allow = causal[None] & same & valid_key
176
+ if opts is not None:
177
+ o = torch.full((len(segs), L), OPT_NONE, device=device)
178
+ for b, op in enumerate(opts):
179
+ o[b, : len(op)] = torch.tensor(op, device=device)
180
+ key_is_option = (o[:, None, :] >= 0)
181
+ query_is_decide = (o[:, :, None] == OPT_DECIDE)
182
+ same_option = o[:, None, :] == o[:, :, None]
183
+ allow = allow & (~key_is_option | query_is_decide | same_option)
184
+ allow = allow | torch.eye(L, dtype=torch.bool, device=device)[None]
185
+ return torch.zeros(len(segs), L, L, dtype=dtype, device=device).masked_fill(~allow, torch.finfo(dtype).min)[:, None]
186
+
187
+
188
+ def rows_of(enc):
189
+ """Split a packed encoding into its state and per-question branch rows.
190
+
191
+ Returns (state_ids, state_pos, rows) with rows[k] = {"ids", "pos", "decide", "opts"}: the branch tokens of question
192
+ k with their (already state-continuing) positions, and the readout offsets *within the branch*. Feeding
193
+ state + rows[k] as one causal row is equivalent to the packed block-causal form for that question, on any
194
+ architecture: the row contains exactly the tokens question k may attend to, in the same positions."""
195
+ seg = enc["seg"]; Ls = seg.count(0)
196
+ rows, start = [], Ls
197
+ for k, (d, oi) in enumerate(zip(enc["decide_idx"], enc["opt_idx"]), start=1):
198
+ end = d + 1 # <decide> is the last token of its branch
199
+ if seg[start] != k or seg[end - 1] != k: raise ValueError("branch layout mismatch")
200
+ rows.append({"ids": enc["ids"][start:end], "pos": enc["pos"][start:end], "decide": d - start, "opts": [o - start for o in oi]})
201
+ start = end
202
+ return enc["ids"][:Ls], enc["pos"][:Ls], rows
203
+
204
+
205
+ class PointerHead(nn.Module):
206
+ def __init__(self, d, dp=256):
207
+ """dp = pointer dimension (head capacity knob)."""
208
+ super().__init__()
209
+ self.q, self.k = nn.Linear(d, dp), nn.Linear(d, dp)
210
+ self.scale = 1 / math.sqrt(dp)
211
+ # calibration: logits are divided by this at inference (eval mode) only. 1.0 = raw. A checkpoint carries the value fitted on
212
+ # its in-distribution development rows (scripts/calibrate_checkpoint.py -> head.pt["temperature"]); training always sees T=1 so
213
+ # a fitted value stays meaningful, and the argmax is unchanged by construction.
214
+ self.temperature = 1.0
215
+
216
+ def forward(self, h_decide, h_opts): # [d], [K,d] -> logits [K]
217
+ z = (self.k(h_opts) @ self.q(h_decide)) * self.scale
218
+ return z if self.training or self.temperature == 1.0 else z / self.temperature
219
+
220
+ def many(self, h_decide, h_opts, owner): # [Q,d], [sum K,d], question of each option [sum K] -> logits [sum K]
221
+ """forward() for many questions at once (serving batches): one pass instead of one per question."""
222
+ z = (self.k(h_opts) * self.q(h_decide)[owner]).sum(-1) * self.scale
223
+ return z if self.training or self.temperature == 1.0 else z / self.temperature
224
+
225
+
226
+ # What a loaded model exposes to kev.serve, kev.predictors and the Space: the scoring interface both DecisionModel (torch)
227
+ # and kev.mlx_model.MLXDecisionModel implement. tests/test_mlx.py checks the MLX class against this list.
228
+ SCORING_INTERFACE = ("encode", "forward", "probs", "probs_and_prefix", "probs_with_prefix", "probs_batch", "eval",
229
+ "head", "backend", "dtype", "device", "hybrid", "option_isolation", "prefix_min_tokens")
230
+
231
+
232
+ EAGER_STATES = 4 # long new states (past the graphed state pass) whose eager prefixes one batched run holds at once: each
233
+ # holds its state's keys, values and DeltaNet states (~0.3 GB for 2,200 tokens on Kev-27B)
234
+
235
+
236
+ def probs_one(model, enc, prefix, keep):
237
+ """-> (probs, prefix to keep) for one request, on any backend: the question rows on the cached prefix on a hit, one
238
+ pass that also returns the prefix when it is to be kept, the plain pass otherwise."""
239
+ if prefix is not None: return model.probs_with_prefix(enc, prefix), prefix
240
+ return model.probs_and_prefix(enc) if keep else (model.probs(enc), None)
241
+
242
+
243
+ class DecisionModel(nn.Module):
244
+ def __init__(self, name, tok, device, lora=None, revision=None, attn=None, head_dim=256, option_isolation=False, special_embeddings=False, lora_targets="all", dtype=torch.float32,
245
+ weights=None, direct_load=False):
246
+ """weights: a full-weight checkpoint directory whose saved backbone replaces the base's (kev.checkpoint's loader rule).
247
+ direct_load: load the backbone straight onto `device` (transformers device_map) instead of staging it in host memory;
248
+ full-weight training on several GPUs in one container needs it (N processes x a 51 GB checkpoint otherwise). Off by
249
+ default: it changes where the rotary buffers are computed, so every other path keeps its bits."""
250
+ super().__init__()
251
+ # backbone only (no vocab head): we never generate text.
252
+ # eager on MPS/CPU (known-good with our float 4D mask); SDPA on CUDA (accepts arbitrary additive masks).
253
+ attn = attn or ("sdpa" if str(device).startswith("cuda") else "eager")
254
+ # dtype: fp32 for training and exact evaluation; bf16 is a serving option for large backbones (8B on a 32 GB Mac)
255
+ load = {"dtype": dtype, "attn_implementation": attn}
256
+ if direct_load: load["device_map"] = {"": torch.cuda.current_device() if device == "cuda" else device} # "cuda": under torchrun, this rank's GPU
257
+ self.lm = AutoModel.from_pretrained(weights, **load) if weights else AutoModelForCausalLM.from_pretrained(name, revision=revision, **load).model
258
+ self.pad_id = pad_id(tok)
259
+ # hybrid backbones (Qwen3.5: Gated DeltaNet layers, recurrent) cannot honour the block-causal mask, so every
260
+ # question runs as its own causal row continuing from the state (rows_of). Attention-only backbones keep the
261
+ # packed form; the two agree to fp32 noise (tests/test_model.py::test_rows_match_packed).
262
+ self.hybrid = is_hybrid(self.lm.config)
263
+ if self.hybrid and option_isolation: raise ValueError("option_isolation needs the packed mask; not available on hybrid backbones")
264
+ self.option_isolation = option_isolation
265
+ if lora:
266
+ from peft import LoraConfig, get_peft_model
267
+ extra = {"trainable_token_indices": {"embed_tokens": [tok.convert_tokens_to_ids(t) for t in SPECIAL]}} if special_embeddings else {}
268
+ targets = {"all": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
269
+ "dense": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], # "all" minus the DeltaNet projections on hybrids (retention ablation)
270
+ "attn": ["q_proj", "k_proj", "v_proj", "o_proj"], "qv": ["q_proj", "v_proj"]}[lora_targets]
271
+ if self.hybrid and lora_targets in ("all", "attn"):
272
+ # Gated DeltaNet projections (transformers 5 names, verified on Qwen3_5TextModel); the mixer's out_proj too
273
+ targets = targets + ["in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj"]
274
+ cfg = LoraConfig(task_type="FEATURE_EXTRACTION", r=lora, lora_alpha=2 * lora, lora_dropout=0.05, target_modules=targets, **extra)
275
+ self.lm = get_peft_model(self.lm, cfg)
276
+ self.head = PointerHead(self.lm.config.hidden_size, dp=head_dim)
277
+ self.device = device
278
+ self.to(device)
279
+
280
+ backend = "torch" # kev.mlx_model.MLXDecisionModel is the other implementation of this scoring interface
281
+ graphs = None # kev.cuda_graphs.CudaGraphs for the serving passes of a hybrid backbone on CUDA (LoadOptions.cuda_graphs)
282
+
283
+ @property
284
+ def prefix_min_tokens(self):
285
+ """kev.serve caches the state prefix from this many state tokens. Attention-only backbones: 384, below which the
286
+ branch-only pass is not faster than one packed pass on MPS (per-op overhead). Hybrid backbones: always, because
287
+ their miss path otherwise recomputes the state once per question (Kev-0.8B bf16 on MPS, 5 questions: 1011 -> 413 ms)."""
288
+ return 0 if self.hybrid else 384
289
+
290
+ @property
291
+ def dtype(self):
292
+ return str(next(self.lm.parameters()).dtype).removeprefix("torch.")
293
+
294
+ def encode(self, tok, rec, **kw):
295
+ """encode() with this model's option-isolation setting; use this from serving/eval code."""
296
+ return encode(tok, rec, option_isolation=self.option_isolation, **kw)
297
+
298
+ def hidden(self, enc):
299
+ return self.hidden_batch([enc])[0, : len(enc["ids"])]
300
+
301
+ SHAPE_BUCKET = int(os.environ.get("KEV_SHAPE_BUCKET", "64")) # MPS: pad the sequence to a multiple of this (per-shape kernel warm-up); 1 disables
302
+
303
+ def _pad_rows(self, rows):
304
+ """Right-pad (ids, pos) token rows into [N, L] id / position tensors and a [N, L] attention mask (1 = real token).
305
+ Pads sit after every real token and are masked keys, so they never change a real token's hidden state (parity
306
+ measured exact). On MPS in eval mode L is rounded up to a SHAPE_BUCKET multiple so kernels are warmed per bucket."""
307
+ L = max(len(ids) for ids, _ in rows)
308
+ if str(self.device) == "mps" and not self.training: L = -(-L // self.SHAPE_BUCKET) * self.SHAPE_BUCKET
309
+ ids = torch.full((len(rows), L), self.pad_id, device=self.device)
310
+ pos = torch.zeros((len(rows), L), dtype=torch.long, device=self.device)
311
+ att = torch.zeros((len(rows), L), dtype=torch.long, device=self.device)
312
+ for i, (rid, rpos) in enumerate(rows):
313
+ ids[i, : len(rid)] = torch.tensor(rid, device=self.device); pos[i, : len(rpos)] = torch.tensor(rpos, device=self.device); att[i, : len(rid)] = 1
314
+ return ids, pos, att
315
+
316
+ def hidden_batch(self, encs):
317
+ """[B, L_max, d] hidden states for a right-padded batch of encoded records under the packed block-causal mask."""
318
+ ids, pos, _ = self._pad_rows([(e["ids"], e["pos"]) for e in encs])
319
+ isolate = any(e.get("option_isolation") for e in encs)
320
+ if isolate and not all(e.get("option_isolation") for e in encs):
321
+ raise ValueError("cannot mix option-isolated and plain encodings in one batch")
322
+ lm_dtype = next(self.lm.parameters()).dtype
323
+ mask = branch_mask_batch([e["seg"] for e in encs], self.device, dtype=lm_dtype, opts=[e["opt"] for e in encs] if isolate else None, length=ids.shape[1])
324
+ return self.lm(input_ids=ids, position_ids=pos, attention_mask=mask).last_hidden_state.float() # head stays fp32
325
+
326
+ def _readout(self, h, enc):
327
+ return [self.head(h[d], h[torch.tensor(oi, device=self.device)]) for d, oi in zip(enc["decide_idx"], enc["opt_idx"])]
328
+
329
+ def rows_form(self, encs):
330
+ """Whether these records run as causal rows: hybrid backbones always (the recurrent layers cannot honour the packed
331
+ mask), attention-only ones when the packed sequence would exceed ROW_PASS_TOKENS (its L x L mask grows without
332
+ bound with the number of questions). The two forms agree (tests/test_model.py::test_rows_match_packed)."""
333
+ return self.hybrid or any(len(e["ids"]) > ROW_PASS_TOKENS for e in encs)
334
+
335
+ def _rows_hidden(self, rows, cache=None, prefix_len=0):
336
+ """Hidden states of causal token rows, one [L_i, d] tensor per row. In eval mode the rows go through the backbone
337
+ rows_per_pass at a time; training keeps one batch (its batches are small and autograd needs the whole graph anyway).
338
+ With `cache`, the rows are branches continuing the cached state: in eager mode the cache is replicated once per
339
+ chunk, leaving the caller's prefix pristine, and the cached tokens are marked real in the attention mask."""
340
+ chunk = len(rows) if self.training else rows_per_pass([ids for ids, _ in rows], prefix_len)
341
+ out = []
342
+ for start in range(0, len(rows), chunk):
343
+ part = rows[start:start + chunk]
344
+ ids, pos, att = self._pad_rows(part)
345
+ past = {}
346
+ if cache is not None:
347
+ replica = copy.copy(cache)
348
+ replica.layers = [copy.copy(layer) for layer in cache.layers]
349
+ for source, target in zip(cache.layers, replica.layers):
350
+ if isinstance(source, LinearAttentionCacheLayerMixin):
351
+ target.conv_states = source.conv_states.copy()
352
+ target.recurrent_states = source.recurrent_states.copy()
353
+ target.is_conv_states_initialized = source.is_conv_states_initialized.copy()
354
+ target.is_recurrent_states_initialized = source.is_recurrent_states_initialized.copy()
355
+ target.has_previous_state = source.has_previous_state.copy()
356
+ target.conv_kernel_size = source.conv_kernel_size.copy()
357
+ replica.reorder_cache(torch.zeros(len(part), dtype=torch.long, device=self.device))
358
+ att = torch.cat([torch.ones((len(part), prefix_len), dtype=torch.long, device=self.device), att], 1)
359
+ past = {"past_key_values": replica, "use_cache": True}
360
+ h = self.lm(input_ids=ids, position_ids=pos, attention_mask=att, **past).last_hidden_state.float()
361
+ out += [h[i, : len(row_ids)] for i, (row_ids, _) in enumerate(part)]
362
+ return out
363
+
364
+ def forward_rows_batch(self, encs, shared_prefix=False):
365
+ """Row form: every question of every record is one causal row = state tokens + its branch tokens. Returns the same
366
+ nested logits as forward_batch. Exact isolation by construction (rows are independent); the state is recomputed
367
+ per row (Q x state tokens), which training accepts; serving and probs() use the prefix cache instead. shared_prefix
368
+ (training, hybrid Qwen3.5): the same rows computed with each state run once and its branches continuing from it,
369
+ gradients included (kev.shared_prefix)."""
370
+ if shared_prefix:
371
+ from .shared_prefix import branch_hidden # here, not at the top: the HF Space vendors model.py alone
372
+ splits = [rows_of(e) for e in encs]
373
+ return [[self.head(h[r["decide"]], h[torch.tensor(r["opts"], device=self.device)]) for h, r in zip(hs, rows)]
374
+ for hs, (_, _, rows) in zip(branch_hidden(self.lm, splits, self.pad_id, self.device), splits)]
375
+ rows, readouts = [], [] # one causal row per question; readouts[i] = (record, <decide> offset, option offsets)
376
+ for b, e in enumerate(encs):
377
+ S, Sp, brs = rows_of(e)
378
+ for r in brs:
379
+ rows.append((S + r["ids"], Sp + r["pos"])); readouts.append((b, len(S) + r["decide"], [len(S) + o for o in r["opts"]]))
380
+ out = [[] for _ in encs]
381
+ for h, (b, d, oi) in zip(self._rows_hidden(rows), readouts):
382
+ out[b].append(self.head(h[d], h[torch.tensor(oi, device=self.device)]))
383
+ return out
384
+
385
+ def forward(self, enc):
386
+ """Returns list of logits tensors, one per question."""
387
+ return self.forward_batch([enc])[0]
388
+
389
+ def forward_batch(self, encs, shared_prefix=False):
390
+ """List (per record) of lists (per question) of logits. Row form or packed block-causal mask, see rows_form
391
+ (the packed mask already runs each state once; shared_prefix does that for the row form of a hybrid backbone)."""
392
+ if self.rows_form(encs): return self.forward_rows_batch(encs, shared_prefix and self.hybrid)
393
+ hs = self.hidden_batch(encs)
394
+ return [self._readout(hs[b], e) for b, e in enumerate(encs)]
395
+
396
+ @torch.no_grad()
397
+ def probs(self, enc):
398
+ """Probabilities per question. A record that runs as rows (every record on a hybrid backbone) takes the serving miss
399
+ path: its state once, then the question rows from its cache. forward() keeps the row form, which re-runs the state
400
+ per question; the two agree to fp32 rounding (#77)."""
401
+ if self.rows_form([enc]): return self.probs_and_prefix(enc)[0]
402
+ return [F.softmax(z, -1).cpu() for z in self.forward(enc)]
403
+
404
+ # --- state-prefix reuse (serving): the state is encoded once, question branches attend to its cached keys/values.
405
+ # Exact by construction: branch tokens never attend to each other across questions (block-causal mask) and the state
406
+ # never sees the branches (causal), so the state's hidden states and KV are identical with or without the branches.
407
+
408
+ def _branch_rows_from_prefix(self, enc, cache):
409
+ """Row-form serving: the branches run as causal rows continuing the cached state (exactly the forward_rows_batch
410
+ layout, minus the recomputed state)."""
411
+ S, _, rows = rows_of(enc)
412
+ hs = self._rows_hidden([(r["ids"], r["pos"]) for r in rows], cache=cache, prefix_len=len(S))
413
+ ps = [F.softmax(self.head(h[r["decide"]], h[torch.tensor(r["opts"], device=self.device)]), -1) for h, r in zip(hs, rows)]
414
+ return list(torch.cat(ps).cpu().split([len(p) for p in ps])) # one device sync for the request, not one per question
415
+
416
+ @torch.no_grad()
417
+ def prefix(self, enc):
418
+ """Run the state tokens only. Returns (n_state_tokens, kv cache, state hidden states [Ls, d] fp32 or None). Only the
419
+ packed pass (probs_with_prefix) reads the hidden states; a hybrid backbone always runs rows, so it keeps None, as the
420
+ graphed prefixes do (a cached fp32 [Ls, d] copy is 160 MiB for 8k tokens on Kev-27B). An attention-only backbone
421
+ keeps them even when this record runs as rows: the same state with fewer questions may be packed."""
422
+ Ls = enc["seg"].count(0)
423
+ ids = torch.tensor([enc["ids"][:Ls]], device=self.device); pos = torch.tensor([enc["pos"][:Ls]], device=self.device)
424
+ # the cache must know the layer types (hybrid backbones keep recurrent + conv states per DeltaNet layer)
425
+ out = self.lm(input_ids=ids, position_ids=pos, past_key_values=DynamicCache(config=self.lm.config), use_cache=True)
426
+ return Ls, out.past_key_values, None if self.hybrid else out.last_hidden_state[0].float()
427
+
428
+ @torch.no_grad()
429
+ def probs_and_prefix(self, enc):
430
+ """One full pass that also returns the state prefix (KV cropped to the state, state hidden states or None, see prefix): a cache miss
431
+ costs a single forward pass, not two."""
432
+ Ls = enc["seg"].count(0)
433
+ if self.rows_form([enc]):
434
+ # recurrent layers cannot be cropped back to the state (and an over-long packed pass is what the row form avoids),
435
+ # so a miss here is a state pass (kept as the prefix) plus the branch rows
436
+ Ls, cache, h_state = self.prefix(enc)
437
+ return self._branch_rows_from_prefix(enc, cache), (Ls, cache, h_state)
438
+ ids = torch.tensor([enc["ids"]], device=self.device); pos = torch.tensor([enc["pos"]], device=self.device)
439
+ dt = next(self.lm.parameters()).dtype
440
+ mask = branch_mask_batch([enc["seg"]], self.device, dtype=dt, opts=[enc["opt"]] if enc.get("option_isolation") else None)
441
+ out = self.lm(input_ids=ids, position_ids=pos, attention_mask=mask, past_key_values=DynamicCache(config=self.lm.config), use_cache=True)
442
+ h = out.last_hidden_state[0].float()
443
+ out.past_key_values.crop(-(len(enc["ids"]) - Ls)) # keep the state only (negative = drop that many trailing tokens; positive form deprecated in transformers 5)
444
+ return [F.softmax(z, -1).cpu() for z in self._readout(h, enc)], (Ls, out.past_key_values, h[:Ls].clone())
445
+
446
+ @torch.no_grad()
447
+ def probs_with_prefix(self, enc, prefix):
448
+ """probs() for a record whose state tokens equal the cached prefix's; only the branches run. The cache is cropped
449
+ back to the state afterwards so it can be reused."""
450
+ Ls, cache, h_state = prefix
451
+ if enc["seg"].count(0) != Ls: raise ValueError("prefix does not match this record's state")
452
+ if self.rows_form([enc]):
453
+ return self._branch_rows_from_prefix(enc, cache)
454
+ ids = torch.tensor([enc["ids"][Ls:]], device=self.device); pos = torch.tensor([enc["pos"][Ls:]], device=self.device)
455
+ dt = next(self.lm.parameters()).dtype
456
+ mask = branch_mask_batch([enc["seg"]], self.device, dtype=dt, opts=[enc["opt"]] if enc.get("option_isolation") else None)[:, :, Ls:, :]
457
+ try:
458
+ out = self.lm(input_ids=ids, position_ids=pos, past_key_values=cache, attention_mask=mask, use_cache=True)
459
+ h = torch.cat([h_state, out.last_hidden_state[0].float()], 0)
460
+ finally:
461
+ cache.crop(-(len(enc["ids"]) - Ls))
462
+ return [F.softmax(z, -1).cpu() for z in self._readout(h, enc)]
463
+
464
+ @torch.no_grad()
465
+ def probs_batch(self, encs, prefixes, keep):
466
+ """kev.serve's model thread: several requests at once. prefixes[i] is request i's cached state prefix or None;
467
+ keep[i] whether to return its new prefix. -> (probs per request, prefix per request: the cached one, the new one if
468
+ kept, else None). With CUDA graphs the requests the graphed passes admit run together (kev.cuda_graphs: shared
469
+ state and row passes; a state too long for the graphed state pass gets its own eager pass first); the rest, and
470
+ every other backend, one at a time. Rows are independent, so a request's answers do not depend on the batch."""
471
+ splits = [rows_of(e) for e in encs]
472
+ def fits(i, cached):
473
+ S, _, rows = splits[i]
474
+ return self.graphs is not None and self.graphs.admits(len(S), [len(r["ids"]) for r in rows], cached)
475
+ long = {i for i in range(len(encs)) if prefixes[i] is None and not fits(i, False) and fits(i, True)}
476
+ batched = [i for i in range(len(encs)) if i in long or fits(i, prefixes[i] is not None)]
477
+ out = {i: probs_one(self, encs[i], prefixes[i], keep[i]) for i in sorted(set(range(len(encs))) - set(batched))}
478
+ runs, held = [[]], 0 # graphed runs, each holding at most EAGER_STATES long states' eager prefixes at once
479
+ for i in batched:
480
+ if i in long and held == EAGER_STATES: runs.append([]); held = 0
481
+ runs[-1].append(i); held += i in long
482
+ from .cuda_graphs import Request # here, not at the top: the HF Space vendors model.py without cuda_graphs.py
483
+ for run in filter(None, runs):
484
+ pre = {i: self.prefix(encs[i]) if i in long else prefixes[i] for i in run}
485
+ X, caches = self.graphs.run([Request(splits[i][0], splits[i][1], [(r["ids"], r["pos"]) for r in splits[i][2]],
486
+ None if pre[i] is None else pre[i][1], keep[i], [[r["decide"], *r["opts"]] for r in splits[i][2]])
487
+ for i in run]) # per question its picks: the <decide> position, then its options
488
+ ps = iter(self._readout_many(X, [len(r["opts"]) for i in run for r in splits[i][2]]))
489
+ for i, cache in zip(run, caches):
490
+ made = pre[i] if i in long else None if cache is None else (len(splits[i][0]), cache, None)
491
+ out[i] = [next(ps) for _ in splits[i][2]], prefixes[i] or (made if keep[i] else None)
492
+ return [out[i][0] for i in range(len(encs))], [out[i][1] for i in range(len(encs))]
493
+
494
+ def _readout_many(self, X, ks):
495
+ """Probabilities of many questions from their picked hidden states X, laid out per question as its <decide> then
496
+ its ks[q] options, in a fixed number of kernels and one device sync. -> one tensor per question."""
497
+ owner = [q for q, k in enumerate(ks) for _ in range(k)]
498
+ slot = [j for k in ks for j in range(k)]
499
+ starts = [0]
500
+ for k in ks: starts.append(starts[-1] + 1 + k)
501
+ dec = torch.tensor(starts[:-1]).to(self.device, non_blocking=True)
502
+ opt, own, sl = torch.tensor([[starts[q] + 1 + j for q, j in zip(owner, slot)], owner, slot]).to(self.device, non_blocking=True)
503
+ z = self.head.many(X[dec], X[opt], own)
504
+ Z = torch.full((len(ks), max(ks)), float("-inf"), device=self.device).index_put_((own, sl), z)
505
+ return torch.softmax(Z, -1)[own, sl].cpu().split(ks)
506
+
507
+ def trainable_parameters(self):
508
+ return [p for p in self.parameters() if p.requires_grad]
runtime/kev/packed_model.py ADDED
@@ -0,0 +1,102 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Standalone MLX mixed ternary/4-bit backbone. Never reads or reconstructs base weights.
2
+
3
+ Ternary codes use MLX's native 2-bit affine kernel; CLoQ stays a separate
4
+ low-rank operation. Only selected embedding rows are dequantized. This is a
5
+ different arithmetic path from the experiment's merged-bf16 simulation.
6
+ """
7
+ import json
8
+ from pathlib import Path
9
+
10
+ import mlx.core as mx
11
+ import mlx.nn as nn
12
+ from mlx.utils import tree_unflatten
13
+ from mlx_lm.models import qwen3_5
14
+
15
+ from .mlx_model import MLXDecisionModel
16
+
17
+ FORMAT = "kev-packed-mlx-v1"
18
+
19
+
20
+ class PackedLinear(nn.QuantizedLinear):
21
+ def __init__(self, input_dims, output_dims, bits, group_size, rank=0, bias=False):
22
+ # Avoid QuantizedLinear's random dense initialization entirely.
23
+ nn.Module.__init__(self)
24
+ if bits not in (2, 4, 8) or input_dims % group_size or input_dims % (32 // bits):
25
+ raise ValueError("unsupported packed linear dimensions or precision")
26
+ self.bits, self.group_size, self.mode = bits, group_size, "affine"
27
+ self.weight = mx.zeros((output_dims, input_dims * bits // 32), dtype=mx.uint32)
28
+ self.scales = mx.zeros((output_dims, input_dims // group_size), dtype=mx.float16)
29
+ self.biases = mx.zeros_like(self.scales)
30
+ if bias:
31
+ self.bias = mx.zeros((output_dims,), dtype=mx.bfloat16)
32
+ if rank:
33
+ self.comp_a = mx.zeros((rank, input_dims), dtype=mx.bfloat16)
34
+ self.comp_b = mx.zeros((output_dims, rank), dtype=mx.bfloat16)
35
+ self.freeze()
36
+
37
+ def __call__(self, x):
38
+ y = super().__call__(x)
39
+ if "comp_a" in self:
40
+ # Keep the small compensation in fp32; never materialize B @ A.
41
+ xf = x.astype(mx.float32)
42
+ y = y + (xf @ self.comp_a.astype(mx.float32).T) @ self.comp_b.astype(mx.float32).T
43
+ return y.astype(mx.bfloat16)
44
+
45
+
46
+ class PackedEmbedding(nn.QuantizedEmbedding):
47
+ activation_dtype = mx.bfloat16
48
+
49
+ def __init__(self, num_embeddings, dims, bits, group_size):
50
+ nn.Module.__init__(self)
51
+ if bits not in (2, 4, 8) or dims % group_size or dims % (32 // bits):
52
+ raise ValueError("unsupported packed embedding dimensions or precision")
53
+ self.num_embeddings, self.dims = num_embeddings, dims
54
+ self.bits, self.group_size, self.mode = bits, group_size, "affine"
55
+ self.weight = mx.zeros((num_embeddings, dims * bits // 32), dtype=mx.uint32)
56
+ self.scales = mx.zeros((num_embeddings, dims // group_size), dtype=mx.float16)
57
+ self.biases = mx.zeros_like(self.scales)
58
+ self.freeze()
59
+
60
+ def __call__(self, x):
61
+ return super().__call__(x).astype(mx.bfloat16)
62
+
63
+
64
+ def load_packed_model(directory, pad_id, head_dim):
65
+ """Construct lazy architecture, replace quantized modules, strictly load packed tensors."""
66
+ root = Path(directory)
67
+ spec = json.loads((root / "packed.json").read_text(encoding="utf-8"))
68
+ if spec.get("format") != FORMAT:
69
+ raise ValueError("unsupported packed model format")
70
+ config = json.loads((root / "config.json").read_text(encoding="utf-8"))
71
+ args = qwen3_5.ModelArgs.from_dict(config)
72
+ text = qwen3_5.TextModelArgs.from_dict(args.text_config)
73
+ pattern = ["full_attention" if (i + 1) % text.full_attention_interval == 0 else "linear_attention" for i in range(text.num_hidden_layers)]
74
+ if args.text_config.get("layer_types", pattern) != pattern:
75
+ raise ValueError("packed layer_types do not match MLX's attention layout")
76
+ # Lazy initialization is discarded before evaluation, including the embedding and LM head.
77
+ lm = qwen3_5.Model(args)
78
+ if "lm_head" in lm.language_model:
79
+ del lm.language_model["lm_head"]
80
+ modules = dict(lm.named_modules())
81
+ replacements = []
82
+ for name, q in spec["linears"].items():
83
+ old = modules.get(name)
84
+ if not isinstance(old, nn.Linear):
85
+ raise ValueError(f"packed linear {name} is not in the architecture")
86
+ out_dims, in_dims = old.weight.shape
87
+ replacements.append((name, PackedLinear(in_dims, out_dims, q["bits"], q["group_size"], q.get("rank", 0), "bias" in old)))
88
+ embedding = spec["embedding"]
89
+ name = "language_model.model.embed_tokens"
90
+ replacements.append((name, PackedEmbedding(text.vocab_size, text.hidden_size, embedding["bits"], embedding["group_size"])))
91
+ lm.update_modules(tree_unflatten(replacements))
92
+ del modules, replacements
93
+ weights = mx.load(str(root / "model.safetensors"))
94
+ if any(v.dtype != mx.uint32 for k, v in weights.items() if k.removesuffix(".weight") in spec["linears"] or k == name + ".weight"):
95
+ raise ValueError("quantized weights must be packed uint32")
96
+ lm.load_weights(list(weights.items()), strict=True)
97
+ del weights
98
+ lm.eval()
99
+ mx.set_cache_limit(64 << 20)
100
+ mx.eval(lm.parameters())
101
+ mx.clear_cache()
102
+ return MLXDecisionModel(lm, pad_id, head_dim)
runtime/kev/plot.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Plot training loss from train logs and held-out accuracy vs baselines from eval.json.
2
+
3
+ Run: uv run python -m kev.plot --logs runs/train.log:kev-0.5b runs/train_holdout.log:holdout --eval runs/kev/eval.json --out docs/training.png
4
+ """
5
+ import argparse, re
6
+ from .suite import read_json
7
+ import matplotlib
8
+ matplotlib.use("Agg")
9
+ import matplotlib.pyplot as plt
10
+
11
+ STEP_RE = re.compile(r"ep(\d+) step (\d+)/(\d+) loss ([\d.]+)")
12
+ ORDER = ["agnews", "agnews_yn", "banking77", "boolq", "mnli", "sst5", "yelp", "yelp_yn"]
13
+ LABEL = {"agnews": "AG News\nK=4", "agnews_yn": "AG News\nyes/no", "banking77": "Banking77\nK=77", "boolq": "BoolQ\nyes/no",
14
+ "mnli": "MNLI\nK=3", "sst5": "SST-5\nscore", "yelp": "Yelp\nscore", "yelp_yn": "Yelp\nyes/no"}
15
+
16
+
17
+ def read_log(path):
18
+ steps, losses, epoch_ends, prev_ep = [], [], [], 0
19
+ for m in STEP_RE.finditer(open(path, encoding="utf-8").read()):
20
+ ep, step, loss = int(m[1]), int(m[2]), float(m[4])
21
+ if ep != prev_ep: epoch_ends.append(steps[-1]); prev_ep = ep
22
+ steps.append(step); losses.append(loss)
23
+ return steps, losses, epoch_ends
24
+
25
+
26
+ def main():
27
+ ap = argparse.ArgumentParser()
28
+ ap.add_argument("--logs", nargs="+", default=["runs/train.log:kev-0.5b"], help="path:label pairs")
29
+ ap.add_argument("--eval", default="runs/kev/eval.json")
30
+ ap.add_argument("--out", default="docs/training.png")
31
+ a = ap.parse_args()
32
+
33
+ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4.2), dpi=160, gridspec_kw={"width_ratios": [1.1, 1.3]})
34
+
35
+ for spec in a.logs:
36
+ path, label = spec.split(":") if ":" in spec else (spec, spec)
37
+ steps, losses, ends = read_log(path)
38
+ ax1.plot(steps, losses, lw=1.6, label=f"{label} (final {losses[-1]:.2f})")
39
+ for e in ends: ax1.axvline(e, color="0.8", lw=0.8, ls="--")
40
+ ax1.set_yscale("log"); ax1.set_xlabel("optimizer step (8 records each)"); ax1.set_ylabel("train loss (cross-entropy, log scale)")
41
+ ax1.set_title("training loss"); ax1.grid(alpha=0.25); ax1.legend(frameon=False)
42
+
43
+ e = read_json(a.eval)
44
+ kev, b0, b1 = e["accuracy_calibration"], e.get("baseline_zero_shot_base", {}), e.get("baseline_zero_shot_instruct", {})
45
+ xs = range(len(ORDER)); w = 0.27
46
+ ax2.bar([x - w for x in xs], [b0.get(s, {}).get("acc", 0) for s in ORDER], w, color="0.78", label="zero-shot base (letter logits)")
47
+ ax2.bar([x for x in xs], [b1.get(s, {}).get("acc", 0) for s in ORDER], w, color="0.55", label="zero-shot Instruct (letter logits)")
48
+ ax2.bar([x + w for x in xs], [kev[s]["acc"] for s in ORDER], w, color="#1f77b4", label="kev-0.5b (pointer readout)")
49
+ for x, s in zip(xs, ORDER):
50
+ ax2.text(x + w, kev[s]["acc"] + 0.012, f"ECE\n{kev[s]['ece']:.2f}", ha="center", va="bottom", fontsize=6.5, color="#1f77b4")
51
+ ax2.set_xticks(list(xs)); ax2.set_xticklabels([LABEL[s] for s in ORDER], fontsize=7.5)
52
+ ax2.set_ylim(0, 1.12); ax2.set_ylabel("held-out accuracy"); ax2.set_title(f"held-out accuracy vs zero-shot baselines (all: {kev['ALL']['acc']:.3f}, ECE {kev['ALL']['ece']:.3f})")
53
+ ax2.grid(axis="y", alpha=0.25); ax2.legend(frameon=False, fontsize=8, loc="upper center", bbox_to_anchor=(0.5, -0.16), ncol=3)
54
+ for ax in (ax1, ax2): ax.spines[["top", "right"]].set_visible(False)
55
+ fig.tight_layout(); fig.savefig(a.out, bbox_inches="tight"); print("wrote", a.out)
56
+
57
+
58
+ if __name__ == "__main__":
59
+ main()
runtime/kev/predictors.py ADDED
@@ -0,0 +1,341 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Predictors: callables record -> {"probabilities": {qid: {key: p}}, "latency_ms", "input_tokens", ...} for kev.benchmark.
2
+
3
+ LocalPredictor scores a checkpoint in-process; RemotePredictor any TypeSafe System One-compatible endpoint; JevPredictor
4
+ Jev itself through the AI SDK worker in playground/scripts (budget-capped).
5
+ """
6
+ import contextlib
7
+ import inspect
8
+ import json
9
+ import math
10
+ import os
11
+ import subprocess
12
+ import sys
13
+ import time
14
+ import urllib.error
15
+ import urllib.request
16
+ from datetime import datetime, timezone
17
+ from importlib.metadata import PackageNotFoundError, version
18
+ from pathlib import Path
19
+
20
+ import torch
21
+ from torch.nn.attention import SDPBackend, sdpa_kernel
22
+
23
+ from kev.api import question_keys
24
+ from kev.checkpoint import Checkpoint, LoadOptions
25
+ from kev.data import api_request, materialize
26
+ from kev.device import sync
27
+ from kev.model import ROW_PASS_TOKENS, ContextOverflow, rows_of
28
+ from kev.suite import CONTEXT
29
+
30
+
31
+ LONG_ROW_KERNELS = "efficient" # the label of a prediction (and its benchmark rows) that took the long-row kernels below
32
+
33
+ KERNEL_PACKAGES = ("torch", "transformers", "peft", "flash-linear-attention", "triton", "causal-conv1d", "mlx", "mlx-lm")
34
+ DELTANET_KERNELS = ("causal_conv1d_fn", "torch_chunk_gated_delta_rule") # the prefill hooks of transformers' GatedDeltaNet
35
+
36
+
37
+ @contextlib.contextmanager
38
+ def long_row_kernels():
39
+ """The attention kernels a long row runs under (LocalPredictor): SDPA's flash and memory-efficient kernels, with the
40
+ math kernel as the fallback. transformers asks SDPA for grouped-query attention (`enable_gqa`) when a call has no
41
+ mask, which is the case for an unpadded state (kev.shared_prefix passes no mask so the state runs causal). Of the
42
+ kernels allowed here only flash and math take `enable_gqa`, and flash has no fp32. So an fp32 state pass used to fall
43
+ back to math, whose L x L scores do not fit: Kev-4B has 16 heads, so at 32k that is 16 x 32k^2 x 4 B = 69 GB per
44
+ layer. Inside this context an fp32 call repeats its keys and values per query head instead (transformers' own path
45
+ whenever a mask is given), and the memory-efficient kernel takes that: fp32 inputs and output, memory linear in L.
46
+ On sm80+ its fp32 products are cutlass's OpMultiplyAddFastF32 (three TF32 products per product, fp32 accuracy and
47
+ fp32 accumulation), not plain TF32. bf16 / fp16 calls keep `enable_gqa` and the flash kernel, so Kev-27B's long rows
48
+ run as before."""
49
+ from transformers.integrations import sdpa_attention
50
+ grouped = sdpa_attention.use_gqa_in_sdpa
51
+ sdpa_attention.use_gqa_in_sdpa = lambda attention_mask, key, value: key.dtype != torch.float32 and grouped(attention_mask, key, value)
52
+ try:
53
+ with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]):
54
+ yield
55
+ finally:
56
+ sdpa_attention.use_gqa_in_sdpa = grouped
57
+
58
+
59
+ def bound_implementation(fn):
60
+ """The function transformers actually calls for one of its kernel hooks: `use_kernel_func_from_hub_with_fallback`
61
+ wraps the PyTorch reference and binds the package's kernel (fla, causal-conv1d) as `implementation` when it imports;
62
+ any other function (unwrapped, or patched in by a caller) is itself. -> "module.qualname"."""
63
+ impl = inspect.getclosurevars(fn).nonlocals.get("implementation", fn) if inspect.isfunction(fn) else fn
64
+ return f"{getattr(impl, '__module__', None)}.{getattr(impl, '__qualname__', type(impl).__name__)}"
65
+
66
+
67
+ def kernel_environment(model, device):
68
+ """What a read's logits depend on besides kev's code and the checkpoint (runs/drift-v1/REPORT.md: #125's causal-conv1d
69
+ kernel moved Kev-27B v1's reads by up to 0.06 in p with no code change): the scoring packages' versions, the GPU, the
70
+ backbone dtype and, on a torch backbone, the attention implementation and (hybrid) the Gated DeltaNet layer's forward
71
+ (kev.fused_qwen35 replaces it) and the convolution and chunked delta rule transformers bound. Reads only what is
72
+ already loaded or imported. Two reads are comparable bit for bit only when this and the kev commit match. Recorded
73
+ in report.json."""
74
+ def installed(name):
75
+ try: return version(name)
76
+ except PackageNotFoundError: return None
77
+ packages = {name: installed(name) for name in KERNEL_PACKAGES}
78
+ packages["torch"] = torch.__version__ # with its build tag (2.8.0+cu128 on the CUDA wheels), which the metadata version drops
79
+ env = {"packages": packages, "device": str(device),
80
+ "gpu": torch.cuda.get_device_name(torch.device(device)) if str(device).startswith("cuda") else None,
81
+ "backend": model.backend, "dtype": model.dtype, "triton_f32_default": os.environ.get("TRITON_F32_DEFAULT"),
82
+ "attention": None, "deltanet": None}
83
+ if model.backend != "torch": return env # kev.mlx_model: mlx-lm's Metal kernels, named by the mlx / mlx-lm versions
84
+ env["attention"] = getattr(model.lm.config, "_attn_implementation", None)
85
+ layer = next((m for m in model.lm.modules() if type(m).__name__.endswith("GatedDeltaNet")), None) if model.hybrid else None
86
+ if layer is not None:
87
+ module = sys.modules[type(layer).__module__]
88
+ env["deltanet"] = {"forward": bound_implementation(getattr(layer.forward, "__func__", layer.forward)),
89
+ **{name: bound_implementation(getattr(module, name)) for name in DELTANET_KERNELS if hasattr(module, name)}}
90
+ return env
91
+
92
+
93
+ class LocalPredictor:
94
+ """Scores a checkpoint in-process. Evaluation is fp32-exact in PyTorch (no TF32, no fused SDPA kernels on CUDA). On
95
+ CUDA a hybrid backbone's Gated DeltaNet layers run flash-linear-attention's Triton kernels, whose dots are TF32, and a
96
+ bf16 backbone is bf16 throughout, so reads repeat bit for bit only on the same kernel set (`self.environment`, from
97
+ kernel_environment; AGENTS.md "What fp32-exact guarantees"). One more exception: on CUDA with the torch backend, a
98
+ record whose longest row (state + one question) exceeds kev.model.ROW_PASS_TOKENS runs under SDPA's flash /
99
+ memory-efficient kernels (long_row_kernels), because the exact math kernel's L x L score matrix does not fit (~275 GB
100
+ per layer pass at 64k for Kev-4B's 16 heads). An fp32 backbone stays fp32 there: the memory-efficient kernel takes
101
+ fp32 (Kev-4B / 9B at 8k-16k against the math kernel: max |dp| 6e-4, no flips in 238 questions;
102
+ scripts/long_state_memory.py). Such a prediction carries `"kernels": LONG_ROW_KERNELS`, kev.benchmark copies it onto
103
+ the record's rows and counts them in report.json's `long_rows`; shorter records carry nothing, so their rows and
104
+ reports are unchanged. The efficient path is CUDA-only: on CPU / MPS attention stays eager and exact and still
105
+ materialises L x L. A long record on a hybrid torch backbone also runs its state once, its questions continuing from
106
+ it (kev.shared_prefix), instead of once per question; the MLX backend's forward already runs the state once."""
107
+
108
+ def __init__(self, run, device, opts=LoadOptions(), context=CONTEXT):
109
+ """opts.temperature=None scores with the temperature the checkpoint carries; 1.0 scores raw logits. context: the
110
+ max_state / max_branch / max_packed a record must encode within (a suite manifest's `context`; the training
111
+ context by default, the serving limits for external suites frozen without admission)."""
112
+ if opts.temperature is not None and not (math.isfinite(opts.temperature) and opts.temperature > 0):
113
+ raise ValueError("temperature must be finite and positive")
114
+ checkpoint = self.checkpoint = Checkpoint(run)
115
+ self.run = checkpoint.path
116
+ if device == "cuda":
117
+ # evaluation is fp32-exact: TF32 (10-bit mantissa) moves probabilities by ~1e-3, the isolation gate's tolerance
118
+ torch.backends.cuda.matmul.allow_tf32 = False; torch.backends.cudnn.allow_tf32 = False
119
+ torch.backends.cuda.enable_flash_sdp(False); torch.backends.cuda.enable_mem_efficient_sdp(False)
120
+ self.tok, self.model = checkpoint.load(device, opts)
121
+ self.temperature = self.model.head.temperature
122
+ self.environment = kernel_environment(self.model, device)
123
+ self.device = device
124
+ self.context = context
125
+
126
+ @torch.no_grad()
127
+ def __call__(self, record):
128
+ enc = self.model.encode(self.tok, materialize(record), max_state=self.context["max_state"], max_branch=self.context["max_branch"], strict=True)
129
+ if len(enc["ids"]) > self.context["max_packed"]:
130
+ raise ContextOverflow(f"packed request exceeds the {self.context['max_packed']}-token limit")
131
+ sync(self.device)
132
+ start = time.perf_counter()
133
+ state, _, rows = rows_of(enc)
134
+ long = len(state) + max(len(r["ids"]) for r in rows) > ROW_PASS_TOKENS
135
+ torch_backend = self.model.backend == "torch"
136
+ efficient = long and torch_backend and self.device == "cuda"
137
+ with long_row_kernels() if efficient else contextlib.nullcontext():
138
+ if long and torch_backend and self.model.hybrid: logits = self.model.forward_batch([enc], shared_prefix=True)[0]
139
+ else: logits = self.model.forward(enc)
140
+ ps = [torch.softmax(z, -1).cpu() for z in logits]
141
+ zs = [z.float().cpu() for z in logits]
142
+ sync(self.device)
143
+ keys = {qid: question_keys(q["type"], q.get("criteria")) for qid, q in record["questions"].items()}
144
+ return {"probabilities": {qid: dict(zip(keys[qid], p.tolist())) for qid, p in zip(keys, ps)},
145
+ "logits": {qid: dict(zip(keys[qid], z.tolist())) for qid, z in zip(keys, zs)},
146
+ "inference_temperature": self.temperature,
147
+ "latency_ms": 1000 * (time.perf_counter() - start), "input_tokens": len(enc["ids"]), **({"kernels": LONG_ROW_KERNELS} if efficient else {})}
148
+
149
+
150
+ REFUSAL_STATUSES = (400, 413, 422) # an endpoint declining the request itself (kev.serve: 422 past its context; Jev: 400 / 413 / 422)
151
+
152
+
153
+ def http_detail(error):
154
+ """The message of an HTTP error body: FastAPI's / TypeSafe's `detail` (or `error`) when it is JSON, else its text."""
155
+ try: text = error.read().decode("utf-8", "replace")
156
+ except Exception: return str(error.reason)
157
+ try: body = json.loads(text)
158
+ except ValueError: return text[:500] or str(error.reason)
159
+ return str(body.get("detail") or body.get("error") or body) if isinstance(body, dict) else str(body)
160
+
161
+
162
+ class RemotePredictor:
163
+ """Score any TypeSafe System One-compatible endpoint (POST <base_url>/v1/systemone) on frozen records. Probabilities are
164
+ taken from the response as returned (renormalised by validate_distribution like every other predictor). Records the
165
+ server-reported model id so the manifest can pin what was scored. `concurrency` is how many requests kev.benchmark may
166
+ keep in flight at once (each call is independent: one request, its own retries); 1 scores sequentially. An endpoint
167
+ that declines the request itself (REFUSAL_STATUSES: kev.serve's 422 for a state past its context) raises
168
+ ContextOverflow at once, so kev.benchmark counts the record rejected under skip_overlong and stops otherwise; another
169
+ client error stops at once too; 408, 429, 5xx, connection errors and timeouts are tried `retries` times in all."""
170
+
171
+ def __init__(self, base_url, model="kev-latest", api_key="local", timeout=120, retries=3, concurrency=1):
172
+ if concurrency < 1:
173
+ raise ValueError("concurrency must be >= 1")
174
+ self.base_url, self.model, self.api_key, self.timeout, self.retries = base_url.rstrip("/"), model, api_key, timeout, retries
175
+ self.concurrency = concurrency
176
+ self.served_model = None
177
+
178
+ def __call__(self, record):
179
+ payload = json.dumps({**api_request(record), "model": self.model}).encode()
180
+ req = urllib.request.Request(f"{self.base_url}/v1/systemone", data=payload, method="POST",
181
+ headers={"content-type": "application/json", "authorization": f"Bearer {self.api_key}"})
182
+ last = None
183
+ for attempt in range(self.retries):
184
+ try:
185
+ start = time.perf_counter()
186
+ with urllib.request.urlopen(req, timeout=self.timeout) as resp:
187
+ body = json.loads(resp.read())
188
+ latency = 1000 * (time.perf_counter() - start)
189
+ break
190
+ except urllib.error.HTTPError as error:
191
+ if error.code in REFUSAL_STATUSES: # the request itself (kev.serve: a state or question over the serving context): a rejected record, not retried
192
+ raise ContextOverflow(f"remote endpoint refused the request (HTTP {error.code}): {http_detail(error)}") from None
193
+ if error.code not in (408, 429) and error.code < 500: # 401, 403, 404, ...: the same answer every time; 408, 429 and 5xx are retried
194
+ raise RuntimeError(f"remote endpoint answered HTTP {error.code}: {http_detail(error)}") from None
195
+ last = error
196
+ except Exception as error: # connection errors and timeouts: retried
197
+ last = error
198
+ if attempt + 1 < self.retries: time.sleep(2 ** attempt)
199
+ else:
200
+ raise RuntimeError(f"remote endpoint failed after {self.retries} attempts: {last}")
201
+ self.served_model = body.get("model", self.served_model)
202
+ probs = {}
203
+ for qid, q in record["questions"].items():
204
+ a = body["answers"][qid]
205
+ if q["type"] == "noul": probs[qid] = {"true": float(a["noul"]), "false": 1 - float(a["noul"])}
206
+ else: probs[qid] = {str(k): float(v) for k, v in a["probabilities"].items()}
207
+ return {"probabilities": probs, "latency_ms": latency, "input_tokens": (body.get("usage") or {}).get("input_tokens")}
208
+
209
+
210
+ PRICE_PER_MILLION = 0.042 # Jev list price per million input tokens, for the budget cap
211
+
212
+
213
+ class RotationAveraged:
214
+ """Test-time order averaging (PLAN.md round 4, item 4.4): score the record under the first `rotations` cyclic
215
+ rotations of every Choice question's options (rotation r of a K-option question is r mod K; Noul and Score keep their
216
+ order, which is part of their meaning) and average the logits per option key. The geometric mean of softmax
217
+ probabilities is the softmax of the mean logits, so the returned probabilities and logits stay consistent; predictors
218
+ that return no logits are averaged in log-probability space. Costs `rotations` forward passes per record."""
219
+
220
+ def __init__(self, predictor, rotations):
221
+ if rotations < 2:
222
+ raise ValueError("rotation averaging needs at least 2 rotations")
223
+ self.predictor, self.rotations = predictor, rotations
224
+ self.temperature = getattr(predictor, "temperature", None)
225
+ self.concurrency = getattr(predictor, "concurrency", 1)
226
+
227
+ @staticmethod
228
+ def rotated(record, r):
229
+ def rotate(q):
230
+ if q["type"] != "choice": return q
231
+ keys = list(q["criteria"]); k = r % len(keys)
232
+ return {**q, "criteria": {key: q["criteria"][key] for key in keys[k:] + keys[:k]}}
233
+ return {**record, "questions": {qid: rotate(q) for qid, q in record["questions"].items()}}
234
+
235
+ def __call__(self, record):
236
+ widest = max((len(q["criteria"]) for q in record["questions"].values() if q["type"] == "choice"), default=1)
237
+ preds = [self.predictor(self.rotated(record, r)) for r in range(min(self.rotations, widest))]
238
+ field = "logits" if all("logits" in p for p in preds) else "probabilities"
239
+ out = {"probabilities": {}, "latency_ms": sum(p["latency_ms"] for p in preds), "rotations": len(preds)}
240
+ if field == "logits":
241
+ out["logits"] = {}; out["inference_temperature"] = preds[0].get("inference_temperature", 1.0)
242
+ for qid in record["questions"]:
243
+ keys = list(preds[0]["probabilities"][qid])
244
+ if field == "logits":
245
+ z = {k: sum(p["logits"][qid][k] for p in preds) / len(preds) for k in keys}
246
+ else:
247
+ z = {k: sum(math.log(max(p["probabilities"][qid][k], 1e-12)) for p in preds) / len(preds) for k in keys}
248
+ top = max(z.values()); e = {k: math.exp(v - top) for k, v in z.items()}; s = sum(e.values())
249
+ out["probabilities"][qid] = {k: v / s for k, v in e.items()}
250
+ if field == "logits": out["logits"][qid] = z
251
+ return out
252
+
253
+
254
+ class JevRefused(ContextOverflow):
255
+ """Jev declined a request because of the input itself (REFUSAL_STATUSES, or a hosted-side error on an oversize request;
256
+ see JevPredictor). With count_refusals, kev.benchmark counts such a record as rejected (rejected.json) instead of stopping."""
257
+
258
+
259
+ # Jev's context in tokens (documented ~32k). Past it the AI Gateway answers with HTTP 400 *or* a 503 GatewayInternalServerError
260
+ # (evals/longdoc-v1: 212 x 400 and 28 x 503 on the 240 requests of ~57k-61k tokens; every request of <= ~31k tokens answered),
261
+ # so a hosted-side error on a request estimated past OVERSIZE x JEV_CONTEXT_TOKENS is the size, not an outage. The estimate is
262
+ # characters / 4 of the request JSON; the 1.5 margin keeps a ~30k-token legal text (~144k characters) out of it.
263
+ JEV_CONTEXT_TOKENS = 32768
264
+ OVERSIZE = 1.5
265
+
266
+
267
+ def oversize(request):
268
+ """Whether a Jev request is well past Jev's context (estimated tokens = characters / 4)."""
269
+ return len(json.dumps(request, ensure_ascii=False)) / 4 > OVERSIZE * JEV_CONTEXT_TOKENS
270
+
271
+
272
+ class JevPredictor:
273
+ """Jev through the AI SDK worker (playground/scripts/jev-evaluate.mjs); every call is counted against a token budget.
274
+ count_refusals: a request answered with a REFUSAL_STATUSES error, or with a hosted-side error (5xx, no status) while
275
+ oversize(), raises JevRefused at once (counted in accounting()["refusals"], not fatal, not retried); any other client error
276
+ (401, 403, 429, ...) still stops the read. attempts: tries per request on hosted-side failures of requests that are not
277
+ oversize, backing off 1, 2, 4, ... seconds up to 30 between them."""
278
+ def __init__(self, key, budget=0.1, max_calls=700, count_refusals=False, attempts=4):
279
+ self.budget, self.max_calls, self.count_refusals, self.attempts = budget, max_calls, count_refusals, attempts
280
+ self.calls, self.input_tokens, self.output_tokens, self.retries = 0, 0, 0, 0
281
+ self.refusals = {} # HTTP status -> count
282
+ self.started_at = datetime.now(timezone.utc).isoformat()
283
+ worker = Path(__file__).resolve().parents[1] / "playground/scripts/jev-evaluate.mjs"
284
+ self.process = subprocess.Popen(["node", str(worker)], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
285
+ stderr=subprocess.DEVNULL, text=True, bufsize=1,
286
+ env={**os.environ, "AI_GATEWAY_API_KEY": key})
287
+
288
+ def close(self):
289
+ self.process.stdin.close()
290
+ try:
291
+ self.process.wait(timeout=5)
292
+ except subprocess.TimeoutExpired:
293
+ self.process.terminate()
294
+ self.process.wait(timeout=5)
295
+
296
+ def __call__(self, record):
297
+ if self.calls >= self.max_calls or (self.input_tokens + 65536) * PRICE_PER_MILLION / 1e6 > self.budget:
298
+ raise RuntimeError("Jev evaluation reached the request/token cost cap")
299
+ request = api_request(record)
300
+ for attempt in range(self.attempts):
301
+ self.process.stdin.write(json.dumps(request) + "\n")
302
+ self.process.stdin.flush()
303
+ line = self.process.stdout.readline()
304
+ if not line:
305
+ raise RuntimeError("Jev SDK worker exited without a response")
306
+ result = json.loads(line)
307
+ self.calls += 1
308
+ if "error" not in result:
309
+ break
310
+ status = result["error"]["status"]
311
+ # bounded retry for hosted-side failures only; client errors (4xx) are real and must surface
312
+ hosted = status is None or status >= 500
313
+ if self.count_refusals and (status in REFUSAL_STATUSES or hosted and oversize(request)):
314
+ kind = str(status) if status in REFUSAL_STATUSES else f"{status} (oversize)"
315
+ self.refusals[kind] = self.refusals.get(kind, 0) + 1
316
+ raise JevRefused(f"Jev refused the request: {result['error']['name']} (HTTP {kind})")
317
+ if attempt == self.attempts - 1 or (status is not None and status < 500):
318
+ raise RuntimeError(f"Jev request failed: {result['error']['name']} (HTTP {status})")
319
+ self.retries += 1
320
+ time.sleep(min(2 ** attempt, 30))
321
+ usage = result["usage"]
322
+ if usage.get("inputTokens") is None:
323
+ raise RuntimeError("Jev returned no input token usage; cannot account for cost")
324
+ self.input_tokens += usage["inputTokens"]
325
+ self.output_tokens += usage.get("outputTokens") or 0
326
+ probabilities = {}
327
+ for qid, q in record["questions"].items():
328
+ answer = result["answers"][qid]
329
+ if q["type"] == "noul":
330
+ probabilities[qid] = {"false": 1 - answer["probability"], "true": answer["probability"]}
331
+ else:
332
+ probabilities[qid] = answer["probabilities"]
333
+ return {**result, "probabilities": probabilities}
334
+
335
+ def accounting(self):
336
+ return {"model": "typesafe-ai/jev", "model_revision": "Gateway alias; provider revision not exposed by SDK result",
337
+ "started_at": self.started_at, "calls": self.calls, "input_tokens": self.input_tokens,
338
+ "output_tokens": self.output_tokens, "retries_after_5xx": self.retries, "listed_input_usd_per_million": PRICE_PER_MILLION,
339
+ **({"refusals": dict(sorted(self.refusals.items()))} if self.count_refusals else {}),
340
+ "estimated_usd": self.input_tokens * PRICE_PER_MILLION / 1e6,
341
+ "budget_usd": self.budget, "sdk": "ai@7.0.105", "zero_data_retention_requested": True}
runtime/kev/publish.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Publish a trained run to the Hugging Face Hub.
2
+
3
+ Repo naming: jaredpalmer/kev-<base size>, e.g. kev-0.5b (Qwen2.5-0.5B), kev-0.6b (Qwen3-0.6B), kev-4b, kev-8b.
4
+ The exact base checkpoint is recorded in the model card's `base_model` field and in head.pt.
5
+
6
+ uv run python -m kev.publish --run runs/kev --repo jaredpalmer/kev-0.5b
7
+ uv run python -m kev.publish --run runs/kev2 --repo jaredpalmer/kev-0.5b --message "v0.2: none-of-the-above fix"
8
+
9
+ Uploads: the adapter (or, for a full-weight run, config.json and every model*.safetensors shard with its index), head.pt,
10
+ tokenizer files, eval.json, training log (if found), and the model card (--card) as README.md with the repo id and run
11
+ name filled in. Requires `hf auth login`. With --private the repo is created private if missing and an existing repo
12
+ that is not private is refused (kev.mirror.ensure_private), so a candidate never lands in a public repo. Full-weight
13
+ shards are linked into the staging directory, not copied (modal_app.py::release_publish uploads a 27B from a CPU container).
14
+
15
+ --replace makes the target revision hold exactly this upload: every file the upload does not carry (other than
16
+ .gitattributes) is deleted in the same commit. A full-weight upload into a repo that still holds an adapter (or an adapter
17
+ upload into one that holds backbone shards) is refused without it, because the loader rule (kev.checkpoint) would pick
18
+ the stale layout.
19
+ """
20
+ import argparse, os, re, shutil, tempfile
21
+ from pathlib import Path
22
+ from huggingface_hub import HfApi
23
+ from .checkpoint import Checkpoint
24
+ from .mirror import ensure_private
25
+ from .suite import read_json, write_json
26
+
27
+ FILES = ["head.pt", "tokenizer.json", "tokenizer_config.json", "chat_template.jinja",
28
+ "vocab.json", "merges.txt", "added_tokens.json", "special_tokens_map.json", "eval.json"]
29
+ WEIGHTS = {False: ["adapter_config.json", "adapter_model.safetensors"], True: ["config.json", "model.safetensors.index.json"]} # by Checkpoint.full
30
+
31
+
32
+ def stale_layout(api, repo, revision, full):
33
+ """Files of the other checkpoint layout on the target revision (adapter files for a full-weight upload, backbone shards
34
+ for an adapter upload), sorted; [] for a missing branch or an empty repo."""
35
+ from huggingface_hub.utils import RevisionNotFoundError
36
+ try: files = api.list_repo_files(repo, repo_type="model", revision=revision)
37
+ except RevisionNotFoundError: return []
38
+ other = (lambda f: f in WEIGHTS[False]) if full else (lambda f: re.fullmatch(r"model.*\.safetensors", f) or f == "model.safetensors.index.json")
39
+ return sorted(f for f in files if other(f))
40
+
41
+
42
+ def main():
43
+ ap = argparse.ArgumentParser()
44
+ ap.add_argument("--run", required=True)
45
+ ap.add_argument("--repo", required=True, help="e.g. jaredpalmer/kev-0.5b")
46
+ ap.add_argument("--message", default=None)
47
+ ap.add_argument("--card", required=True, help="model card markdown, e.g. docs/model-cards/kev-4b.md")
48
+ ap.add_argument("--private", action="store_true")
49
+ ap.add_argument("--tag", help="create this Hub tag on the uploaded commit (versioned release, e.g. v0.2)")
50
+ ap.add_argument("--revision", help="upload to this branch instead of main (created if missing); for candidates that must not replace the released weights")
51
+ ap.add_argument("--replace", action="store_true", help="delete every file of the target revision this upload does not carry, in the same commit (tag the previous version first)")
52
+ a = ap.parse_args()
53
+
54
+ checkpoint, run_name = Checkpoint(a.run), os.path.basename(a.run.rstrip("/"))
55
+ base, full = checkpoint.meta.base, checkpoint.full
56
+ api = HfApi()
57
+ if a.private: ensure_private(api, a.repo) # creates a missing repo private; refuses one that exists public
58
+ else: api.create_repo(a.repo, repo_type="model", exist_ok=True, private=False)
59
+ stale = stale_layout(api, a.repo, a.revision, full)
60
+ if stale and not a.replace:
61
+ raise SystemExit(f"{a.repo}{'@' + a.revision if a.revision else ''} holds {stale}, the other checkpoint layout; the loader would pick it. "
62
+ "Tag the previous version, then publish with --replace to delete it in the same commit")
63
+
64
+ with tempfile.TemporaryDirectory() as tmp:
65
+ shards = {p.name for p in checkpoint.shards()} if full else set()
66
+ for f in FILES + WEIGHTS[full] + sorted(shards):
67
+ src = f"{a.run}/{f}"
68
+ if not os.path.exists(src): print(f"skip {f} (not found)")
69
+ elif f in shards: os.symlink(os.path.abspath(src), f"{tmp}/{f}") # read-only and ~5 GB each: linked, not copied (upload_folder follows links)
70
+ else: shutil.copy(src, tmp)
71
+ if not full: # runs before task_type was set saved null; the Hub warns about it and PEFT treats both the same for a bare backbone
72
+ cfg_path = f"{tmp}/adapter_config.json"; cfg = read_json(cfg_path)
73
+ if not cfg.get("task_type"): cfg["task_type"] = "FEATURE_EXTRACTION"; write_json(cfg_path, cfg)
74
+ if os.path.exists(f"runs/logs/train_{run_name}.log"): # standalone runs keep their log there (see .gitignore)
75
+ shutil.copy(f"runs/logs/train_{run_name}.log", f"{tmp}/train.log")
76
+ # research trials: runs/<study>/<trial>/checkpoint -> ship the trial's result, provenance and training log too
77
+ trial = os.path.dirname(a.run.rstrip("/")) if run_name == "checkpoint" else None
78
+ if trial:
79
+ run_name = os.path.relpath(trial, "runs")
80
+ for f in ("result.json", "provenance.json", "train.log", "training_config.json", "training_metrics.json", "interpolation.json"):
81
+ for src in (f"{trial}/{f}", f"{a.run}/{f}"):
82
+ if os.path.exists(src): shutil.copy(src, f"{tmp}/{f}"); break
83
+
84
+ # the card's prose names the Hub repo and trial itself; only the frontmatter is filled from the checkpoint
85
+ card = re.sub(r"^base_model: .*$", f"base_model: {base}", Path(a.card).read_text(encoding="utf-8"), flags=re.M)
86
+ if "base_model_relation:" not in card: card = card.replace(f"base_model: {base}", f"base_model: {base}\nbase_model_relation: {'finetune' if full else 'adapter'}")
87
+ Path(tmp, "README.md").write_text(card, encoding="utf-8")
88
+
89
+ ev = read_json(f"{tmp}/eval.json") if os.path.exists(f"{tmp}/eval.json") else {}
90
+ acc = ev.get("accuracy_calibration", {}).get("ALL", {})
91
+ if os.path.exists(f"{tmp}/result.json"):
92
+ clean = read_json(f"{tmp}/result.json").get("clean", {}) # research trials; other result files (kev-finetune runs) just skip the figures
93
+ acc = {"acc": clean.get("acc", float("nan")), "ece": clean.get("ece", float("nan"))}
94
+ msg = a.message or f"Upload {run_name} (base {base}; acc {acc.get('acc', float('nan')):.3f}, ECE {acc.get('ece', float('nan')):.3f})"
95
+ if a.revision: api.create_branch(a.repo, branch=a.revision, repo_type="model", exist_ok=True)
96
+ # delete_patterns="*": remote files not in this upload are deleted in the same commit (huggingface_hub keeps .gitattributes and every file being uploaded)
97
+ info = api.upload_folder(folder_path=tmp, repo_id=a.repo, repo_type="model", commit_message=msg, revision=a.revision,
98
+ delete_patterns="*" if a.replace else None)
99
+ print(info)
100
+ if a.tag:
101
+ api.create_tag(a.repo, tag=a.tag, repo_type="model", tag_message=msg, exist_ok=False, revision=a.revision)
102
+ print(f"tagged {a.repo}@{a.tag}")
103
+
104
+
105
+ if __name__ == "__main__":
106
+ main()
runtime/kev/rounds.py ADDED
@@ -0,0 +1,1364 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Registered research rounds as data: one spec per round, one engine for every stage.
2
+
3
+ A round (PLAN.md, rounds 5-20) always has the same shape. Arms are config-only trials (or checkpoints made from them,
4
+ round 20's interpolations) compared with a parent checkpoint;
5
+ each arm and its parent are read on a list of suites; a rule compares every arm with its parent on paired,
6
+ record-clustered bootstraps (primaries with a lower bound, guards with thresholds, pooled panels); the passing arm with
7
+ the best score is the size's candidate; confirmation stages (test partitions, the locked read) are read once. The spec
8
+ (`experiments/rounds/r<N>.json`, schema below) is committed before any training or read, like the PLAN registration it
9
+ encodes.
10
+
11
+ uv run python -m kev.rounds validate experiments/rounds/r18.json # well-formed, every parent read exists, plans validate
12
+ uv run python -m kev.rounds launch experiments/rounds/r18.json # modal_app.py::study for each study (after `modal deploy`)
13
+ uv run python -m kev.rounds watch experiments/rounds/r18.json # pull + read each finished trial once, then the read-out
14
+ uv run python -m kev.rounds launch-reads experiments/rounds/r18.json [--arms a,b] [--stage tests --arm a]
15
+ uv run python -m kev.rounds readout experiments/rounds/r18.json # -> runs/r18-readout/round18.json + a table
16
+ uv run python -m kev.rounds confirm experiments/rounds/r18.json --stage tests [--arm a] # -> runs/r18-verdict/<size>-<stage>.json
17
+
18
+ Every comparison is served-vs-served: each side's rows are served at the temperature fitted on its own trial's
19
+ development rows (`kev.metrics.served`), or, for the arms of a round that registers a `temperature` pool, on that pool;
20
+ unknowable records are scored only by `unknowable_report`, and every delta is `kev.metrics.paired_bootstrap` (2,000
21
+ resamples, seed 0, micro), the registered read since round 5.
22
+
23
+ Spec (paths are relative to the repo root; templates take {round}, {arm}, {size}, {tag}):
24
+ round, registered the round number and where its registration lives
25
+ archive a recorded round: the git tag holding the plans, runs and data it names (validate lists
26
+ what this checkout lacks instead of failing; launch refuses a record)
27
+ gpu, app Modal GPU and KEV_APP_NAME for launches (optional)
28
+ studies {name: {plan, suite, transfer, gpu, timeout, budget}} one modal_app.py::study call each (budget >= its admission bound)
29
+ reads {tag: {suite, flags?} | {entrypoint: "locked_test", decision}} what each read tag scores (entrypoint: ENTRYPOINTS)
30
+ read_timeout {size: seconds} overrides modal_app.READ_TIMEOUTS for one size (a 27B's fp32 reads)
31
+ locked_args {size: [args]} extra modal_app.py::locked_test switches for one size (a 27B's GPU memory)
32
+ parents {name: {trial, checkpoint?, reads: {tag: dir}}} checkpoint (Hub id[@rev]) only when /<trial>/checkpoint is not on the volume
33
+ arms {name: {trial?, checkpoint?, parent, reads?, select?, transfer_read?, trained_on?}} name = "<size>-<label>"; reads
34
+ default to arm_reads; select false = reported, never the candidate (attribution arms); an
35
+ arm without a trial (an interpolated checkpoint) names its /runs/... checkpoint, needs the
36
+ round's `temperature` pool (it has no development rows) and, if the rule reads
37
+ "transfer", a transfer_read; trained_on [suite dirs] names what such a checkpoint was
38
+ trained on (default: covered by the round's trial arms, e.g. the trials it interpolates)
39
+ arm_reads template of an arm's read directory (default "runs/r{round}-{arm}-{tag}")
40
+ transfer_read a read tag (round level, or per arm, which wins; null = the trial's own) whose rows are an arm's "transfer"
41
+ rows instead of its trial's in-trial transfer read (a checkpoint without a trial has none)
42
+ temperature {reads: [tags], sources?: {tag: [source]}, exclude_reads?: [tags]} every ARM (parents keep their trial's
43
+ development rows) is served at the temperature fitted, as everywhere (kev.metrics.served),
44
+ on the knowable rows of `reads` pooled (a read listed in `sources` gives only those
45
+ sources' rows), minus every record whose id is in the `exclude_reads` rows; the arm's own
46
+ rule-stage reads are used at every stage; no pooled tag may sit in a panel that a
47
+ temperature-dependent criterion reads (see pool_problems), and no pooled read may share
48
+ data with any arm's training data (pool_training_problems: round 19's failure mode)
49
+ drop_ids record ids dropped on both sides of every comparison (devtools-v1's duplicated ids)
50
+ rule {panels, unknowable?, criteria, rank}
51
+ confirm {stage: {candidate_reads, parent_reads?, panels, criteria}}
52
+ A panel is {reads: [tags], metrics: [bootstrapped], report: [value only], source?: filter, versus?: {name: dir}, by_length?,
53
+ exclude_sources?, exclude_tasks?, exclude_file?, optional?}; the tag "transfer" is the trial's own in-trial transfer read (or
54
+ the arm's transfer_read). source keeps one source (or a list of them); exclude_sources / exclude_tasks drop rows by source /
55
+ task, and exclude_file ({path, sha256}: a JSON object of `ids` (question or record ids), `sources` and `tasks` lists, or a
56
+ bare list of any of them) drops what it lists, checked against its registered hash, so a private list (round 24's
57
+ tasksource-heldout families) enters a public spec only as a path and a sha256 (panel_filter; the same rows leave both
58
+ sides and any `versus` reference; the panel records how many questions they removed as `excluded`). optional: true marks
59
+ a report-only panel no criterion or rank may read; its absent reads leave the comparison complete (`missing_optional`).
60
+ The `temperature` pool may carry ci: {level, samples, seed}, which adds each arm's bootstrap interval of its pooled
61
+ temperature to the read-out (`temperature_ci`, report only). by_length (true, or {edges?:
62
+ [tokens], tokenizer?: [name, revision]}) adds acc / ece / brier / confident_error_rate per state-token bucket
63
+ (kev.metrics.calibration_by_length: "<metric>_<bucket>" entries, e.g. ece_16k_plus, value only) with state tokens counted
64
+ from the reads' suite records. A criterion is {left, op, right, plus?} where left is a path
65
+ ("<panel>.<metric>.<candidate|parent|delta|lower|upper>", "<panel>.<metric>_<bucket>.<candidate|parent>",
66
+ "unknowable.<candidate|parent>") or a list [a, b] meaning a - b, and right is a number or a path (plus is added to it).
67
+ rank is a list of {by: path | [paths summed], order}.
68
+
69
+ Every arm's read-out records where its temperature came from (`temperature_source`: the pool with its reads and question
70
+ count, or its trial's development rows with their suite) and its parent's (`parent_temperature_source`: its trial's
71
+ development rows, plus `shipped`, the head.pt temperature when it is here), and the printed table warns when an arm's rows
72
+ are a training corpus's (in distribution: round 19's failure mode). From round 21 (POOL_REQUIRED_FROM) a rule or
73
+ confirmation with a temperature-dependent criterion must register a `temperature` pool: validate and launch refuse it
74
+ otherwise; earlier rounds get a warning, so their recorded specs keep validating.
75
+
76
+ Removed suites (kev.suite.REMOVED_SUITES: evals/external/scienthoon-v1, removed 2026-09-27, last read by round 22;
77
+ evals/external/wanli-v2 and typesafe-v1, removed 2026-09-30, last read by round 26; wanli-v1, last read by round 5). A round
78
+ up to the suite's `last_round` that reads one validates with the read listed under `archived`; its read-out and
79
+ confirmations are computed from the committed rows as before (they need rows, not the suite). A later round that names
80
+ one is refused with the reason, and no read of one is launched again (read_commands skips it; kev.suite.load_split refuses it).
81
+ """
82
+ import argparse
83
+ import functools
84
+ import operator
85
+ import os
86
+ import socket
87
+ import subprocess
88
+ import sys
89
+ import time
90
+ from pathlib import Path
91
+ from typing import NamedTuple
92
+
93
+ from kev.metrics import (LENGTH_EDGES, LENGTH_METRICS, calibration_by_length, length_buckets, metrics, paired_bootstrap, raw_row, recorded,
94
+ scored_rows, served, served_at, tempered_row, unknowable_report)
95
+ from kev.suite import ADMISSION_TOKENIZER, digest, file_lock, load_split, read_json, read_manifest, removed_suite, suite_key, write_json
96
+
97
+ ROOT = Path(__file__).resolve().parents[1]
98
+ SAMPLES = 2000 # the registered resample count since round 5
99
+ OPS = {">": operator.gt, ">=": operator.ge, "<": operator.lt, "<=": operator.le}
100
+ FIELDS = ("candidate", "parent", "delta", "lower", "upper")
101
+ TRANSFER = "transfer" # the in-trial transfer read (runs/<study>/<trial>/transfer)
102
+ STAGGER = 60 # seconds between `modal run` launches: Modal refuses more than ~3 app creations a minute
103
+ ENTRYPOINTS = ("benchmarks", "locked_test") # the modal_app.py entrypoints a read can go through
104
+ IN_FLIGHT = 7200 # seconds an arm's launched reads are presumed running (the longest suite timeout in modal_app.READ_TIMEOUTS)
105
+
106
+
107
+ # --- rows ------------------------------------------------------------------------------------------------------------
108
+
109
+ def paired(candidate, reference, metric):
110
+ """The registered paired read: record-clustered bootstrap of candidate - reference, micro, 2,000 resamples, seed 0."""
111
+ x = paired_bootstrap(candidate, reference, samples=SAMPLES, seed=0, metric=metric, aggregation="micro")
112
+ return {"delta": x[f"micro_{metric}_delta"], "ci95": x["ci95"]}
113
+
114
+
115
+ def temperature(trial, root=ROOT):
116
+ """The temperature fitted on a trial's own development rows (decision-v7's, or its suite's; how every checkpoint here is
117
+ served unless the round registers a `temperature` pool for its arms)."""
118
+ return served(read_json(Path(root) / trial / "development/rows.json"), [])[0]
119
+
120
+
121
+ def served_clean(rows, t):
122
+ """Every clean row served at t, unknowable records included (what unknowable_report scores; served_at drops them)."""
123
+ return [tempered_row(raw_row(recorded(r)), t) for r in rows if r["variant"] == "clean"]
124
+
125
+
126
+ class Pool(NamedTuple):
127
+ """Where a pooled temperature is fitted (spec `temperature`, resolved to one arm's read directories)."""
128
+ reads: list # directories whose rows are pooled
129
+ sources: dict # {directory: [source, ...]}: the only sources pooled from that read (absent = every source)
130
+ exclude: list # directories whose record ids are removed from the pool
131
+ ci: dict = None # {level, samples, seed}: also bootstrap the fitted temperature's interval (report only; round 24)
132
+
133
+
134
+ def select_rows(reads, exclude=()):
135
+ """(pooled rows, rows before the exclusion): every rows.json of reads, as (path, sources or None), concatenated with
136
+ each read limited to its sources, minus every record whose id appears in an exclude rows.json. The one row selection
137
+ of a pooled temperature, here and in scripts/calibrate_checkpoint.py (which writes it into a checkpoint)."""
138
+ excluded = {r["id"] for path in exclude for r in read_json(path)}
139
+ rows = [r for path, sources in reads for r in read_json(path) if sources is None or r["source"] in sources]
140
+ return [r for r in rows if r["id"] not in excluded], rows
141
+
142
+
143
+ def pooled_temperature(pool, root=ROOT):
144
+ """(temperature, fit report): the temperature() objective (kev.metrics.served: knowable clean rows, TEMPERATURE_FIT)
145
+ on select_rows of the pool."""
146
+ kept, rows = select_rows([(Path(root) / d / "rows.json", pool.sources.get(d)) for d in pool.reads], [Path(root) / d / "rows.json" for d in pool.exclude])
147
+ return served(kept, [])[0], {"reads": pool.reads, **({"sources": pool.sources} if pool.sources else {}), "exclude_reads": pool.exclude,
148
+ "questions": len(scored_rows(kept)), "excluded_questions": len(scored_rows(rows)) - len(scored_rows(kept))}
149
+
150
+
151
+ def pooled_temperature_ci(pool, root=ROOT):
152
+ """The pooled temperature's bootstrap interval (report only; the audit of 2026-09-27: the served T of a 648-question pool
153
+ moves by ~±0.15, enough to move an ECE criterion by more than its margin): the same fit (TEMPERATURE_FIT's grid and
154
+ objective, knowable clean rows of select_rows) on `samples` resamples of the pool's (source, group) clusters within each
155
+ source (kev.metrics.cluster_resamples, seed), percentile interval at `level`. Each question's NLL on the grid is computed
156
+ once; a resample's fit is the grid point with the least mean NLL over its questions."""
157
+ import numpy as np
158
+ from kev.metrics import TEMPERATURE_FIT, cluster_resamples, nll_at_temperature
159
+ kept, _ = select_rows([(Path(root) / d / "rows.json", pool.sources.get(d)) for d in pool.reads], [Path(root) / d / "rows.json" for d in pool.exclude])
160
+ raw = served_at(kept, 1.0)
161
+ grid = np.exp(np.linspace(np.log(0.25), np.log(4), TEMPERATURE_FIT["points"]))
162
+ nll = np.asarray([[nll_at_temperature(r, float(t)) for t in grid] for r in raw])
163
+ level, samples, seed = pool.ci.get("level", 0.9), pool.ci.get("samples", SAMPLES), pool.ci.get("seed", 0)
164
+ fits = np.asarray([grid[int(np.argmin(nll[idx].mean(axis=0)))] for idx in cluster_resamples(raw, samples, seed)])
165
+ lower, upper = np.quantile(fits, [(1 - level) / 2, 1 - (1 - level) / 2])
166
+ return {"level": level, "samples": samples, "seed": seed, "lower": float(lower), "upper": float(upper), "questions": len(raw),
167
+ "resampled": "(source, group) clusters within each source (kev.metrics.cluster_resamples)"}
168
+
169
+
170
+ class Side:
171
+ """One checkpoint of a comparison: its trial, where each read tag's rows live and its served temperature (fitted on
172
+ the trial's development rows, or on `pool` when the round registers one for its arms)."""
173
+
174
+ def __init__(self, trial, dirs, root=ROOT, drop=(), pool=None, checkpoint=None, suite=None, shipped=None):
175
+ self.trial, self.dirs, self.root, self.drop, self.pool, self.checkpoint = trial, dirs, Path(root), set(drop), pool, checkpoint
176
+ self.suite = suite # the suite the trial trained and was scored on (its development rows'), when known
177
+ self.shipped = shipped # a parent's: () -> the temperature its checkpoint ships (head.pt), or None when not here
178
+ self._t, self._served, self.fit = None, {}, None
179
+
180
+ def rows_path(self, tag):
181
+ return self.root / self.dirs[tag] / "rows.json"
182
+
183
+ def has(self, tag):
184
+ return tag in self.dirs and self.rows_path(tag).exists()
185
+
186
+ def absent_fit(self):
187
+ """The rows the temperature needs that are not here (pool directories, or the trial's development rows)."""
188
+ needed = [*self.pool.reads, *self.pool.exclude] if self.pool else [f"{self.trial}/development"]
189
+ return [d for d in needed if not (self.root / d / "rows.json").exists()]
190
+
191
+ @property
192
+ def t(self):
193
+ if self._t is None:
194
+ if self.pool: self._t, self.fit = pooled_temperature(self.pool, self.root)
195
+ else: self._t = temperature(self.trial, self.root)
196
+ return self._t
197
+
198
+ def temperature_source(self):
199
+ """Where the served temperature was fitted (read-out provenance): the pool (its reads and knowable question count),
200
+ or the trial's development rows, their suite and whether that suite is a training corpus (then the temperature is
201
+ in distribution: round 19's failure mode)."""
202
+ if self.pool: return {"kind": "pool", "reads": self.pool.reads, "questions": self.fit["questions"] if self.fit else None}
203
+ m = suite_manifest(self.suite) if self.suite else None
204
+ return {"kind": "trial development rows", "rows": f"{self.trial}/development", "suite": self.suite,
205
+ "training_corpus": bool(trainable_sources(m)) if m is not None else None, **({"shipped": self.shipped()} if self.shipped else {})}
206
+
207
+ def served(self, tag):
208
+ if tag not in self._served: self._served[tag] = served_at(read_json(self.rows_path(tag)), self.t)
209
+ return self._served[tag]
210
+
211
+ def panel(self, spec):
212
+ """A panel's served, knowable rows, reads concatenated in the order listed, minus drop_ids and whatever the
213
+ panel's filters leave out (panel_filter: `source`, `exclude_sources`, `exclude_tasks`, `exclude_file`)."""
214
+ keep = panel_filter(spec, self.root)
215
+ return [r for tag in spec["reads"] for r in self.served(tag) if r["id"] not in self.drop and keep(r)]
216
+
217
+ def unknowable_share(self, tag):
218
+ return unknowable_report(served_clean(read_json(self.rows_path(tag)), self.t))["share_at_0_9"] if self.has(tag) else None
219
+
220
+
221
+ # --- panel filters ---------------------------------------------------------------------------------------------------
222
+
223
+ EXCLUDE_KEYS = ("exclude_sources", "exclude_tasks", "exclude_file") # round 24: the audit's per-suite fixes, applied to both sides
224
+ EXCLUDE_FILE_KEYS = ("ids", "sources", "tasks") # what an exclude file may list (other keys, e.g. a salt, are ignored)
225
+
226
+
227
+ def exclude_file_path(entry, root=ROOT):
228
+ return Path(root) / entry["path"]
229
+
230
+
231
+ @functools.cache
232
+ def _exclusion(path, sha256):
233
+ if digest(path) != sha256:
234
+ raise ValueError(f"exclude_file {path} does not match the sha256 the spec registers ({sha256[:12]}...): a changed exclusion list is a new spec")
235
+ data = read_json(path)
236
+ if isinstance(data, list): return {k: frozenset(data) for k in EXCLUDE_FILE_KEYS} # a bare list: ids, sources or tasks alike
237
+ return {k: frozenset(data.get(k, ())) for k in EXCLUDE_FILE_KEYS}
238
+
239
+
240
+ def exclusion(entry, root=ROOT):
241
+ """{ids, sources, tasks} of a panel's exclude_file ({path, sha256}: a JSON object with any of those lists, or a bare list
242
+ matched against all three), checked against its registered sha256. The file may be private (never in git: e.g. round 24's
243
+ tasksource-heldout families, restored with scripts/private_rows.py); the spec carries only its path and hash."""
244
+ return _exclusion(exclude_file_path(entry, root), entry["sha256"])
245
+
246
+
247
+ def panel_filter(panel, root=ROOT):
248
+ """row -> kept? for a panel: `source` (a name, or a list of names) keeps only those sources; `exclude_sources` and
249
+ `exclude_tasks` drop rows by source and task; `exclude_file` drops rows whose id, record (group), source or task it lists.
250
+ The same filter serves both sides of a comparison (and a `versus` reference), so paired deltas stay paired."""
251
+ source = panel.get("source")
252
+ only = None if source is None else {source} if isinstance(source, str) else set(source)
253
+ sources, tasks = set(panel.get("exclude_sources", ())), set(panel.get("exclude_tasks", ()))
254
+ listed = exclusion(panel["exclude_file"], root) if panel.get("exclude_file") else None
255
+ def keep(r):
256
+ if only is not None and r["source"] not in only: return False
257
+ if r["source"] in sources or r.get("task") in tasks: return False
258
+ return not (listed and (r["id"] in listed["ids"] or r.get("group") in listed["ids"] or r["source"] in listed["sources"] or r.get("task") in listed["tasks"]))
259
+ return keep
260
+
261
+
262
+ def filter_problems(panel, spec, root=ROOT, rows=True):
263
+ """What is wrong with a panel's filters: malformed values, an exclude_file without {path, sha256}, and (rows: the local
264
+ data a read-out needs) an exclude_file here that does not match its hash. `exclude_sources` names are checked
265
+ against the manifests of the panel's reads (a typo would silently keep a source the rule means to drop)."""
266
+ problems, source = [], panel.get("source") # an absent exclude_file is not a problem here: validate lists it with the absent rows
267
+ if source is not None and not (isinstance(source, str) or (isinstance(source, list) and source and all(isinstance(s, str) for s in source))):
268
+ problems.append("source is a source name or a non-empty list of names")
269
+ for key in ("exclude_sources", "exclude_tasks"):
270
+ if key in panel and not (isinstance(panel[key], list) and panel[key] and all(isinstance(s, str) for s in panel[key])):
271
+ problems.append(f"{key} is a non-empty list of names")
272
+ if "optional" in panel and not isinstance(panel["optional"], bool): problems.append("optional is true or false")
273
+ entry = panel.get("exclude_file")
274
+ if entry is not None:
275
+ if not (isinstance(entry, dict) and isinstance(entry.get("path"), str) and isinstance(entry.get("sha256"), str) and len(entry["sha256"]) == 64):
276
+ problems.append("exclude_file is {path, sha256}")
277
+ elif rows and exclude_file_path(entry, root).exists() and digest(exclude_file_path(entry, root)) != entry["sha256"]:
278
+ problems.append(f"exclude_file {entry['path']} does not match its registered sha256")
279
+ if isinstance(panel.get("exclude_sources"), list):
280
+ suites = [spec["reads"][t].get("suite") for t in panel["reads"] if t in spec["reads"]]
281
+ if TRANSFER in panel["reads"]:
282
+ tag = spec.get("transfer_read")
283
+ suites.append(spec["reads"].get(tag, {}).get("suite") if tag else None)
284
+ manifests = [suite_manifest(suite_dir(s)) if s and suite_dir(s) else None for s in suites]
285
+ if manifests and all(m is not None for m in manifests):
286
+ listed = set().union(*(listed_sources(m) for m in manifests))
287
+ if unknown := sorted(set(panel["exclude_sources"]) - listed): problems.append(f"exclude_sources name {unknown}, which no read of the panel lists")
288
+ return problems
289
+
290
+
291
+ def gating_panels(rule):
292
+ """Panels a criterion or the rank reads: an `optional` panel (report only) must not be one of them."""
293
+ paths = [p for c in rule["criteria"].values() for p in _paths(c)]
294
+ paths += [p for key in rule.get("rank", []) for p in (key["by"] if isinstance(key["by"], list) else [key["by"]])]
295
+ return {p.split(".")[0] for p in paths}
296
+
297
+
298
+ # --- spec ------------------------------------------------------------------------------------------------------------
299
+
300
+ def load(path):
301
+ spec = read_json(path)
302
+ spec.setdefault("arm_reads", "runs/r{round}-{arm}-{tag}")
303
+ return spec
304
+
305
+
306
+ def size_of(arm):
307
+ return arm.split("-")[0]
308
+
309
+
310
+ def locations(where, spec, **values):
311
+ """{tag: dir} from a template string (one for every tag of the spec) or an explicit {tag: template} mapping."""
312
+ tags = [*spec.get("reads", {}), TRANSFER]
313
+ if isinstance(where, str):
314
+ return {tag: where.format(round=spec["round"], tag=tag, **values) for tag in tags}
315
+ return {tag: d.format(round=spec["round"], tag=tag, **values) for tag, d in where.items()}
316
+
317
+
318
+ def transfer_read(spec, arm):
319
+ """The read tag whose rows are the arm's "transfer" rows, or None for its trial's in-trial transfer read."""
320
+ a = spec["arms"][arm]
321
+ return a["transfer_read"] if "transfer_read" in a else spec.get("transfer_read")
322
+
323
+
324
+ def arm_side(spec, arm, root=ROOT, stage=None):
325
+ """The arm's side. Its "transfer" rows and its temperature pool come from its rule-stage reads at every stage (the
326
+ development reads it was selected on); a stage's candidate_reads locate the stage's own panels."""
327
+ a = spec["arms"][arm]
328
+ development = locations(a.get("reads", spec["arm_reads"]), spec, arm=arm, size=size_of(arm))
329
+ where = spec["confirm"][stage]["candidate_reads"] if stage else a.get("reads", spec["arm_reads"])
330
+ tag = transfer_read(spec, arm)
331
+ transfer = development.get(tag) if tag else f"{a['trial']}/transfer" if a.get("trial") else None
332
+ dirs = {**locations(where, spec, arm=arm, size=size_of(arm)), **({TRANSFER: transfer} if transfer else {})}
333
+ pool, registered = None, spec.get("temperature")
334
+ if registered:
335
+ fit = {**development, **({TRANSFER: transfer} if transfer else {})}
336
+ pool = Pool([fit[r] for r in registered["reads"]], {fit[r]: s for r, s in registered.get("sources", {}).items()}, [fit[r] for r in registered.get("exclude_reads", [])], registered.get("ci"))
337
+ suite = trial_suite(spec, a["trial"], root) if a.get("trial") and not pool else None
338
+ return Side(a.get("trial"), dirs, root, spec.get("drop_ids", ()), pool, a.get("checkpoint"), suite)
339
+
340
+
341
+ def parent_side(spec, arm, root=ROOT, stage=None):
342
+ p = spec["parents"][spec["arms"][arm]["parent"]]
343
+ where = (spec["confirm"][stage].get("parent_reads") if stage else None) or p["reads"]
344
+ dirs = {**locations(where, spec, arm=arm, size=size_of(arm)), TRANSFER: f"{p['trial']}/transfer"}
345
+ return Side(p["trial"], dirs, root, spec.get("drop_ids", ()), suite=trial_suite(spec, p["trial"], root), shipped=lambda: shipped_temperature(p, root))
346
+
347
+
348
+ def rule_tags(rule):
349
+ return {tag for panel in rule["panels"].values() for tag in panel["reads"]} | ({rule["unknowable"]} if rule.get("unknowable") else set())
350
+
351
+
352
+ def _paths(criterion):
353
+ left = criterion["left"] if isinstance(criterion["left"], list) else [criterion["left"]]
354
+ return [*left, *([criterion["right"]] if isinstance(criterion["right"], str) else [])]
355
+
356
+
357
+ class Validation(NamedTuple):
358
+ problems: list # the spec is malformed, or something a launch or a read-out needs is missing
359
+ archived: list # files a recorded round names that this checkout does not carry (they live on spec["archive"])
360
+
361
+
362
+ def validate(spec, root=ROOT, rows=True, plans=True, partitions=False):
363
+ """Validation(problems, archived). rows: every parent's development rows and every read the rule needs from it exist
364
+ under root (what a read-out, and so a launch, needs); plans: every read suite exists, every plan trial passes
365
+ kev.experiment.validated_trial against its suite manifest and fits its budget; partitions: plans go through load_plan
366
+ instead (verifies the partitions, may fetch them from the Hub). A recorded round (`archive`: the git tag holding its
367
+ evidence) lists absent files under `archived` instead of failing; a round without it must have them all. A read of a
368
+ removed suite (kev.suite.REMOVED_SUITES) is archived (with plans) for a round up to the suite's last_round, whether or
369
+ not the round carries `archive`, and a problem for any later round."""
370
+ problems, archived, root = [], [], Path(root)
371
+ absent = (archived if spec.get("archive") else problems).append
372
+ for key in ("round", "registered", "parents", "arms", "reads", "rule"):
373
+ if key not in spec: problems.append(f"missing key {key!r}")
374
+ if problems: return Validation(problems, archived)
375
+ tags = {*spec["reads"], TRANSFER}
376
+ for name, r in spec["reads"].items():
377
+ if r.get("entrypoint", "benchmarks") not in ENTRYPOINTS: problems.append(f"read {name}: entrypoint {r['entrypoint']!r} is not one of {ENTRYPOINTS}"); continue
378
+ if r.get("entrypoint", "benchmarks") == "benchmarks" and not r.get("suite"): problems.append(f"read {name}: no suite")
379
+ if r.get("entrypoint") == "locked_test" and not r.get("decision"): problems.append(f"read {name}: locked_test needs a decision suite")
380
+ if r.get("suite") and (gone := removed_suite(r["suite"])): # kev.suite.REMOVED_SUITES: archived up to its last round, refused after
381
+ if isinstance(spec["round"], int) and spec["round"] <= gone["last_round"]:
382
+ if plans: archived.append(f"read {name}: {r['suite']} not in this checkout: removed on {gone['removed']} (kev.suite.REMOVED_SUITES); "
383
+ f"round {spec['round']} read it before the removal, and the rows it made are the record")
384
+ else:
385
+ problems.append(f"read {name}: {r['suite']} was removed on {gone['removed']} and may not be read after round {gone['last_round']}: {gone['reason']}")
386
+ continue
387
+ if plans and r.get("suite") and not (root / r["suite"]).exists(): absent(f"read {name}: {r['suite']} not in this checkout")
388
+ stages = {None: spec["rule"], **spec.get("confirm", {})}
389
+ reads_transfer = any(TRANSFER in panel["reads"] for rule in stages.values() for panel in rule["panels"].values())
390
+ for name, a in spec["arms"].items():
391
+ if a.get("parent") not in spec["parents"]: problems.append(f"arm {name}: unknown parent {a.get('parent')!r}")
392
+ if "-" not in name: problems.append(f"arm {name}: name must be <size>-<label>")
393
+ tag = transfer_read(spec, name)
394
+ if tag is not None and (tag not in spec["reads"] or spec["reads"][tag].get("entrypoint", "benchmarks") != "benchmarks"):
395
+ problems.append(f"arm {name}: transfer_read {tag!r} is not a benchmarks read of this spec")
396
+ if tag is not None and isinstance(a.get("reads"), dict) and tag not in a["reads"]: problems.append(f"arm {name}: its reads do not locate its transfer_read {tag!r}")
397
+ if not a.get("trial"):
398
+ if not str(a.get("checkpoint", "")).startswith("/runs/"): problems.append(f"arm {name}: an arm without a trial needs a checkpoint on the runs volume (/runs/...)")
399
+ if not spec.get("temperature"): problems.append(f"arm {name}: no trial, so no development rows to fit its temperature on; register a `temperature` pool")
400
+ if reads_transfer and tag is None: problems.append(f"arm {name}: no trial, so no in-trial transfer read; give it a transfer_read")
401
+ problems += pool_problems(spec, stages)
402
+ conflicts, unknown_training = pool_training_problems(spec, root)
403
+ problems += conflicts + pool_required_problems(spec)
404
+ for line in unknown_training: absent(line)
405
+ for stage, rule in stages.items():
406
+ where = f"confirm.{stage}" if stage else "rule"
407
+ for pname, panel in rule["panels"].items():
408
+ unknown = set(panel["reads"]) - tags
409
+ if unknown: problems.append(f"{where} panel {pname}: unknown read tags {sorted(unknown)}")
410
+ if not panel.get("metrics") and not panel.get("report"): problems.append(f"{where} panel {pname}: nothing to compute")
411
+ problems += [f"{where} panel {pname}: {p}" for p in by_length_problems(panel, spec)]
412
+ problems += [f"{where} panel {pname}: {p}" for p in filter_problems(panel, spec, root, rows)]
413
+ if panel.get("optional") and pname in gating_panels(rule): problems.append(f"{where} panel {pname}: optional (report only), but a criterion or the rank reads it")
414
+ for cname, c in rule["criteria"].items():
415
+ if c.get("op") not in OPS: problems.append(f"{where} criterion {cname}: op {c.get('op')!r}")
416
+ for path in _paths(c):
417
+ if not _known_path(path, rule): problems.append(f"{where} criterion {cname}: unknown path {path!r}")
418
+ for key in rule.get("rank", []):
419
+ for path in key["by"] if isinstance(key["by"], list) else [key["by"]]:
420
+ if not _known_path(path, rule): problems.append(f"{where} rank: unknown path {path!r}")
421
+ if stage and "candidate_reads" not in rule: problems.append(f"{where}: no candidate_reads")
422
+ if rule_tags(spec["rule"]) - tags: problems.append(f"rule.unknowable: unknown read tag {spec['rule']['unknowable']!r}")
423
+ for study, s in spec.get("studies", {}).items() if plans else ():
424
+ if not (root / s["plan"]).exists(): absent(f"study {study}: plan {s['plan']} not in this checkout")
425
+ elif (root / s["suite"] / "manifest.json").exists(): problems += _validate_plan(study, s, root, partitions)
426
+ else: absent(f"study {study}: suite {s['suite']} not in this checkout")
427
+ if rows:
428
+ for where, pname, entry in [(s or "rule", n, p["exclude_file"]) for s, rule in stages.items() for n, p in rule["panels"].items() if isinstance(p.get("exclude_file"), dict)]:
429
+ if isinstance(entry.get("path"), str) and not exclude_file_path(entry, root).exists():
430
+ absent(f"{where} panel {pname}: exclude_file {entry['path']} not in this checkout (a private list: scripts/private_rows.py restore)")
431
+ for pname, p in spec["parents"].items():
432
+ if not (root / p["trial"] / "development/rows.json").exists(): absent(f"parent {pname}: {p['trial']}/development/rows.json not in this checkout")
433
+ for arm in spec["arms"]:
434
+ side = parent_side(spec, arm, root)
435
+ required = {t for panel in spec["rule"]["panels"].values() if not panel.get("optional") for t in panel["reads"]} # an optional panel's reads may be absent
436
+ for tag in sorted(required - {spec["rule"].get("unknowable")}): # the parent's unknowable read is reported, not required
437
+ if not side.has(tag): absent(f"arm {arm}: parent read {tag} not in this checkout ({side.dirs.get(tag, 'no location')})")
438
+ return Validation(problems, archived)
439
+
440
+
441
+ def _known_path(path, rule):
442
+ parts = path.split(".")
443
+ if parts[0] == "unknowable": return len(parts) == 2 and parts[1] in ("candidate", "parent") and bool(rule.get("unknowable"))
444
+ panel = rule["panels"].get(parts[0])
445
+ if panel is None or len(parts) != 3 or parts[2] not in FIELDS: return False
446
+ value_only = parts[1] in panel.get("report", ()) or parts[1] in by_length_names(panel) # report and by-length metrics have no interval
447
+ return parts[1] in panel.get("metrics", ()) or (value_only and parts[2] in ("candidate", "parent"))
448
+
449
+
450
+ # --- calibration by state length -------------------------------------------------------------------------------------
451
+
452
+ def by_length_options(panel):
453
+ """(edges, tokenizer) of a panel's by_length option: true = kev.metrics.LENGTH_EDGES and kev.suite.ADMISSION_TOKENIZER."""
454
+ option = panel.get("by_length")
455
+ option = {} if option is True else option
456
+ return tuple(option.get("edges", LENGTH_EDGES)), tuple(option.get("tokenizer", ADMISSION_TOKENIZER))
457
+
458
+
459
+ def by_length_names(panel):
460
+ """{"<metric>_<bucket>": metric} for the entries a by_length panel reports (empty for any other panel, or invalid edges)."""
461
+ if not panel.get("by_length"): return {}
462
+ try:
463
+ return {f"{m}_{name}": m for m in LENGTH_METRICS for name, _, _ in length_buckets(by_length_options(panel)[0])}
464
+ except (ValueError, AttributeError, TypeError):
465
+ return {}
466
+
467
+
468
+ def by_length_problems(panel, spec):
469
+ option = panel.get("by_length")
470
+ if not option: return []
471
+ if option is not True and (not isinstance(option, dict) or set(option) - {"edges", "tokenizer"}):
472
+ return ["by_length is true or {edges?, tokenizer?}"]
473
+ problems = []
474
+ try:
475
+ edges, tokenizer = by_length_options(panel) # tuple() of a non-iterable raises: a malformed option is a problem, not a crash
476
+ length_buckets(edges)
477
+ if len(tokenizer) != 2 or not all(isinstance(x, str) for x in tokenizer): problems.append("by_length: tokenizer is [name, revision]")
478
+ except (ValueError, TypeError) as error:
479
+ problems.append(f"by_length: {error}" if isinstance(error, ValueError) else f"by_length: edges are a list of token counts and tokenizer is [name, revision] ({error})")
480
+ if TRANSFER in panel["reads"]: problems.append("by_length counts state tokens from a read's suite; 'transfer' names no single suite")
481
+ problems += [f"by_length: read {tag!r} has no suite" for tag in panel["reads"] if tag in spec["reads"] and not spec["reads"][tag].get("suite")]
482
+ return problems
483
+
484
+
485
+ @functools.cache
486
+ def state_lengths(suite, split, tokenizer=ADMISSION_TOKENIZER):
487
+ """{record id: state tokens} of a suite partition with the tokenizer (name, revision): the counts a by_length panel
488
+ buckets both sides by. A state's token count is its encoded state segment, as kev.model.encode builds it: the <state>
489
+ token plus user_tokens of the materialised state (len + 1, what encode checks against max_state and kev.serve reports
490
+ as `state_tokens`, seg.count(0)); the one definition behind a row's `state_tokens` (kev.metrics.calibration_by_length)."""
491
+ from kev.data import materialize
492
+ from kev.model import load_tokenizer, user_tokens
493
+ tok = load_tokenizer(*tokenizer)
494
+ return {r["_meta"]["id"]: 1 + len(user_tokens(tok, materialize(r)["state"])) for r in load_split(ROOT / suite, split, allow_test=split == "test")}
495
+
496
+
497
+ def panel_lengths(spec, panel):
498
+ """{record id: state tokens} over a by_length panel's reads (each read's suite and partition)."""
499
+ tokenizer, out = by_length_options(panel)[1], {}
500
+ for tag in panel["reads"]:
501
+ read = spec["reads"][tag]
502
+ for rid, n in state_lengths(read["suite"], read_split(read), tokenizer).items():
503
+ if out.setdefault(rid, n) != n: raise ValueError(f"record {rid!r} has different states in the reads of one by_length panel")
504
+ return out
505
+
506
+
507
+ # --- temperature pools against training data (round 19's failure mode) ------------------------------------------------
508
+
509
+ ROUND_19 = ("round 19's failure mode: a temperature fitted on held-out items of a checkpoint's own training corpus is in distribution "
510
+ "(round 19 served its SFT arms at T 0.955 fitted on sft-v1 development rows: breadth-v1 ECE 0.059); fit on a pool of "
511
+ "held-out datasets instead (round 20: breadth-v1 ECE 0.0085)")
512
+ FIT_SPLITS = ("calibration", "development") # a training corpus's held-out partitions: never a temperature pool
513
+
514
+
515
+ def read_split(read):
516
+ """The partition a read scores: kev.benchmark's development, `--split X`, or the test under --allow-test / locked_test."""
517
+ flags = read.get("flags", "").split()
518
+ if read.get("entrypoint") == "locked_test" or "--allow-test" in flags: return "test"
519
+ for i, flag in enumerate(flags):
520
+ if flag.startswith("--split="): return flag.split("=", 1)[1]
521
+ if flag == "--split" and i + 1 < len(flags): return flags[i + 1]
522
+ return "development"
523
+
524
+
525
+ def suite_dir(path):
526
+ """'evals/...' for a suite path given relative, absolute or as a container saw it (/root/kev/evals/sft-v1); None if none."""
527
+ return suite_key(path)
528
+
529
+
530
+ @functools.cache
531
+ def _suites():
532
+ """({manifest sha256: suite dir}, {directory name: [suite dirs]}) of every suite in this checkout."""
533
+ by_digest, by_name = {}, {}
534
+ for m in sorted((ROOT / "evals").rglob("manifest.json")):
535
+ d = str(m.parent.relative_to(ROOT))
536
+ by_digest[digest(m)] = d
537
+ by_name.setdefault(m.parent.name, []).append(d)
538
+ return by_digest, by_name
539
+
540
+
541
+ def suite_by_digest(sha256):
542
+ """The suite dir whose manifest.json hashes to sha256 (what provenance, head.pt and report.json record), or None."""
543
+ return _suites()[0].get(sha256)
544
+
545
+
546
+ def suite_manifest(suite):
547
+ return read_manifest(ROOT / suite) if (ROOT / suite / "manifest.json").exists() else None
548
+
549
+
550
+ def listed_sources(manifest):
551
+ """Every source a manifest names: trainable, held-out, eval-only and its `sources` table (a dict or a list)."""
552
+ return {*manifest.get("trainable_sources", []), *manifest.get("holdout_sources", []), *manifest.get("eval_only_sources", []), *manifest.get("sources", [])}
553
+
554
+
555
+ def trainable_sources(manifest):
556
+ """The sources a suite offers for training (`trainable_sources`, or `sources` entries marked trainable): a suite with
557
+ any is a training corpus, whose calibration and development partitions are held-out items of training sources."""
558
+ table = manifest.get("sources", {})
559
+ return {*manifest.get("trainable_sources", []), *(k for k, v in table.items() if isinstance(v, dict) and v.get("trainable"))} if isinstance(table, dict) else set(manifest.get("trainable_sources", []))
560
+
561
+
562
+ class Training(NamedTuple):
563
+ """What a checkpoint was trained on, as far as this checkout can tell."""
564
+ suites: frozenset # suite dirs: its training suite, the components that suite names, a plan's `data` directory
565
+ sources: frozenset # the sources those suites train
566
+ unlisted: tuple # training suites whose sources this checkout cannot list
567
+
568
+
569
+ def training_data(suites, data=None):
570
+ """Training over suite dirs plus a `data` file: each suite's trainable sources (every listed source for a data-only
571
+ manifest without `trainable_sources`, e.g. round6/b1v2) plus the suites its manifest names as components (sft-v1's
572
+ `inputs.components`: documents-v1-train -> evals/documents-v1, hard-v1-extra-train -> evals/hard-v1). A `data` file
573
+ counts through its directory's suite; one outside evals/ (a `kev.train --data` run) cannot be checked, so it is
574
+ `unlisted` like any training this checkout cannot place, never silently dropped."""
575
+ data_dir = suite_dir(Path(data).parent) if data else None
576
+ seen, sources, unlisted, todo = set(), set(), [] if data_dir or not data else [f"data {data} (outside evals/)"], [d for d in (*suites, data_dir) if d]
577
+ while todo:
578
+ d = todo.pop(0)
579
+ if d in seen: continue
580
+ seen.add(d)
581
+ m = suite_manifest(d)
582
+ own = (trainable_sources(m) if "trainable_sources" in m else listed_sources(m)) if m is not None else set()
583
+ if not own: unlisted.append(d)
584
+ sources |= own
585
+ inputs = (m or {}).get("inputs")
586
+ for component in inputs.get("components", {}) if isinstance(inputs, dict) else ():
587
+ found = _suites()[1].get(component.removesuffix("-train").removesuffix("-extra"))
588
+ if found: todo += found
589
+ else: unlisted.append(f"{d} component {component}")
590
+ return Training(frozenset(seen), frozenset(sources), tuple(unlisted))
591
+
592
+
593
+ def recorded_training(suite_sha256=None, suite=None, data=None):
594
+ """Training of a checkpoint from what its trainer recorded (provenance.json, head.pt): the suite whose manifest hashes to
595
+ suite_sha256 (else the suite path it was given) and its `data` file's directory; None when no suite of this checkout matches."""
596
+ trained = suite_by_digest(suite_sha256)
597
+ if trained is None and suite and suite_dir(suite) and suite_manifest(suite_dir(suite)) is not None: trained = suite_dir(suite)
598
+ return training_data([trained], data) if trained else None
599
+
600
+
601
+ def _study(spec, trial):
602
+ return next((s for name, s in spec.get("studies", {}).items() if trial.startswith(f"runs/{name}/")), None)
603
+
604
+
605
+ def trial_suite(spec, trial, root=ROOT):
606
+ """The suite a trial trained on and was scored on (its development rows'): its study's in this spec, else the manifest
607
+ its provenance.json hashes; None when neither is here."""
608
+ if study := _study(spec, trial): return suite_dir(study["suite"])
609
+ provenance = Path(root) / trial / "provenance.json"
610
+ return suite_by_digest(read_json(provenance).get("suite_sha256")) if provenance.exists() else None
611
+
612
+
613
+ def trial_training(spec, trial, root=ROOT):
614
+ """Training of a trial: its study's suite plus its plan entry's `data` (trial <NN>-<label> is entry NN), else what its
615
+ provenance.json records; None when neither is here."""
616
+ study, provenance = _study(spec, trial), Path(root) / trial / "provenance.json"
617
+ if study:
618
+ plan = next((p / study["plan"] for p in (Path(root), ROOT) if (p / study["plan"]).exists()), None)
619
+ entries, index = (read_json(plan) if plan else []), Path(trial).name.split("-")[0]
620
+ data = entries[int(index)].get("data") if isinstance(entries, list) and index.isdigit() and int(index) < len(entries) else None
621
+ if plan or not provenance.exists(): return training_data([suite_dir(study["suite"])], data)
622
+ if not provenance.exists(): return None
623
+ p = read_json(provenance)
624
+ return recorded_training(p.get("suite_sha256"), data=p.get("config", {}).get("data"))
625
+
626
+
627
+ def pool_conflicts(pooled, training):
628
+ """[(label, why)] for every way a pooled read shares data with a checkpoint's Training. pooled = [(label, suite, split,
629
+ sources or None)]: (a) the read's suite is training data of the checkpoint (its training suite, or a component or `data`
630
+ suite of it); (b) the sources it pools (its allowlist, else every source its manifest lists) are sources the checkpoint
631
+ trained on; (c) it scores the calibration or development partition of a training corpus (any suite with trainable
632
+ sources). A read whose suite or sources cannot be established is refused too. The one check, for kev.rounds pools and
633
+ scripts/calibrate_checkpoint.py alike."""
634
+ out = []
635
+ for label, suite, split, sources in pooled:
636
+ d = suite_dir(suite) if suite else None
637
+ m = suite_manifest(d) if d else None
638
+ if m is None:
639
+ out.append((label, f"its suite ({suite!r}) has no manifest in this checkout, so what it pools cannot be checked")); continue
640
+ if d in training.suites: out.append((label, f"{d} is training data of the checkpoint"))
641
+ pooled_sources = set(sources) if sources else listed_sources(m)
642
+ if not pooled_sources: out.append((label, f"{d} lists no sources; give the read a `sources` allowlist"))
643
+ if shared := sorted(pooled_sources & training.sources):
644
+ out.append((label, f"it pools {len(shared)} training source(s) {shared[:8]}" + (" ..." if len(shared) > 8 else "")))
645
+ if split in FIT_SPLITS and trainable_sources(m): out.append((label, f"it reads the {split} partition of {d}, a training corpus"))
646
+ return out
647
+
648
+
649
+ def pool_training_problems(spec, root=ROOT):
650
+ """(problems, unknown) of the round's `temperature` pool against every arm's training data (pool_conflicts; message
651
+ names round 19's failure mode). An arm's training comes from `trained_on`, else from its trial (trial_training); an arm
652
+ without either (a checkpoint made from the round's trials) is covered by the trial arms. unknown: arms whose training
653
+ this checkout cannot establish or list (a trial without provenance, a manifest without sources, a `data` file outside
654
+ evals/): a problem for a new round, reported under archived for a recorded one."""
655
+ pool = spec.get("temperature")
656
+ if not pool or not pool.get("reads"): return [], []
657
+ pooled = [(tag, spec["reads"][tag].get("suite"), read_split(spec["reads"][tag]), pool.get("sources", {}).get(tag)) for tag in pool["reads"] if tag in spec["reads"]]
658
+ found, problems, unknown = {}, [], []
659
+ for arm, a in spec["arms"].items():
660
+ if "trained_on" in a:
661
+ if not isinstance(a["trained_on"], list) or not a["trained_on"] or not all(isinstance(s, str) for s in a["trained_on"]):
662
+ problems.append(f"arm {arm}: trained_on is a non-empty list of suite directories"); continue
663
+ training = training_data([suite_dir(s) or s for s in a["trained_on"]])
664
+ elif a.get("trial"):
665
+ training = trial_training(spec, a["trial"], root)
666
+ if training is None:
667
+ unknown.append(f"arm {arm}: what {a['trial']} trained on is unknown (no study of this spec, no provenance.json), so the temperature pool cannot be checked against it"); continue
668
+ else:
669
+ continue
670
+ unknown += [f"arm {arm}: cannot list the sources of {d}, part of its training, so the temperature pool cannot be checked against it" for d in training.unlisted]
671
+ by_tag = {}
672
+ for tag, why in pool_conflicts(pooled, training): by_tag.setdefault(tag, []).append(why)
673
+ for tag, whys in by_tag.items(): found.setdefault((tag, "; ".join(whys)), []).append(arm)
674
+ for (tag, why), arms in found.items():
675
+ problems.append(f"temperature: pooled read {tag!r} shares data with the training of arm(s) {', '.join(arms)}: {why}. This is {ROUND_19}")
676
+ if not any("trained_on" in a or a.get("trial") for a in spec["arms"].values()):
677
+ problems.append("temperature: no arm names its training (a trial or `trained_on`), so the pool cannot be checked against it")
678
+ return problems, unknown
679
+
680
+
681
+ POOL_REQUIRED_FROM = 21 # from this round on, a temperature-dependent criterion without a `temperature` pool is refused
682
+
683
+
684
+ def temperature_criteria(spec):
685
+ """["rule <name>" | "confirm.<stage> <name>"] of every criterion that reads a number the temperature moves."""
686
+ stages = {None: spec["rule"], **spec.get("confirm", {})}
687
+ return sorted({f"{'confirm.' + s if s else 'rule'} {name}" for s, rule in stages.items() for name, c in rule["criteria"].items()
688
+ if any(_moves_with_temperature(path, rule) for path in _paths(c))})
689
+
690
+
691
+ def _listed(names):
692
+ return f"{len(names)} ({', '.join(names[:3])}{', ...' if len(names) > 3 else ''})"
693
+
694
+
695
+ def pool_required_problems(spec):
696
+ """Round >= POOL_REQUIRED_FROM with a temperature-dependent criterion and no `temperature` pool: refused (its arms would
697
+ be served at their trials' development rows, round 19's failure mode). Earlier rounds only warn (calibration_warnings),
698
+ so their recorded specs keep validating."""
699
+ number, moved = spec.get("round"), temperature_criteria(spec)
700
+ if spec.get("temperature") or not moved or not isinstance(number, int) or number < POOL_REQUIRED_FROM: return []
701
+ return [f"temperature: round {number} registers no `temperature` pool, but {_listed(moved)} criteria depend on the temperature, "
702
+ f"so its arms would be served at their trials' development rows. This is {ROUND_19}; register a pool of held-out datasets (copy r20)"]
703
+
704
+
705
+ def shipped_temperature(entry, root=ROOT):
706
+ """The temperature a parent's checkpoint ships (head.pt), read locally only: <trial>/checkpoint/head.pt, else its Hub
707
+ checkpoint's head.pt in the local Hugging Face cache (no network); None when neither is here."""
708
+ path = Path(root) / entry["trial"] / "checkpoint" / "head.pt"
709
+ if not path.exists() and entry.get("checkpoint") and not str(entry["checkpoint"]).startswith("/"):
710
+ from huggingface_hub import try_to_load_from_cache
711
+ repo, _, revision = entry["checkpoint"].partition("@")
712
+ cached = try_to_load_from_cache(repo, "head.pt", revision=revision or None)
713
+ path = Path(cached) if isinstance(cached, str) else path
714
+ if not path.exists(): return None
715
+ from kev.checkpoint import read_meta
716
+ return read_meta(path.parent).temperature
717
+
718
+
719
+ SHIPPED_TOLERANCE = 0.05 # a parent served this far from its shipped temperature, fitted in distribution, is warned about
720
+
721
+
722
+ def continuation_warnings(spec, root=ROOT):
723
+ """Report only: a full-weight study runs longer than one attempt only while something continues it. Its trials spawn
724
+ with Modal's retries off (kev.budget.FULL_FT_RETRIES says why), so a timed-out trial gets its next attempt from
725
+ `kev.rounds watch` (modal_app.continue_full_trial), or by hand with modal_app.py::resume --trial; an attached launch
726
+ keeps no attempt ledger and is never continued."""
727
+ out = []
728
+ for name, study in spec.get("studies", {}).items():
729
+ plan = Path(root) / study.get("plan", "")
730
+ trials = read_json(plan) if plan.is_file() else []
731
+ if any(isinstance(t, dict) and t.get("full_ft") for t in (trials if isinstance(trials, list) else [])):
732
+ out.append(f"study {name} is full-weight: a trial that times out is continued only while `kev.rounds watch` runs (or by hand: "
733
+ f"modal_app.py::resume --study {name} --trial <label>); an attached launch (study --detached False) keeps no ledger and is never continued")
734
+ return out
735
+
736
+
737
+ def calibration_warnings(spec, root=ROOT):
738
+ """Report-only warnings that `validate` and `launch` print. (1) For a round before POOL_REQUIRED_FROM without a
739
+ `temperature` pool whose criteria depend on the temperature (from then on it is a problem: pool_required_problems): each
740
+ trial arm whose development rows are a training corpus's is served in distribution, round 19's failure mode. (2) Every
741
+ parent served at a temperature fitted on a training corpus's development rows that differs from its shipped head.pt
742
+ temperature by more than SHIPPED_TOLERANCE (when both are here)."""
743
+ out, moved = [], temperature_criteria(spec)
744
+ early = isinstance(spec.get("round"), int) and spec["round"] < POOL_REQUIRED_FROM
745
+ for arm, a in spec["arms"].items() if moved and early and not spec.get("temperature") else ():
746
+ suite = trial_suite(spec, a["trial"], root) if a.get("trial") else None
747
+ if suite and trainable_sources(suite_manifest(suite) or {}):
748
+ out.append(f"arm {arm} will be served at a temperature fitted on its trial's {suite} development rows, a training corpus, and "
749
+ f"{_listed(moved)} criteria depend on it: {ROUND_19}; register a `temperature` pool")
750
+ for name, p in spec["parents"].items():
751
+ suite = trial_suite(spec, p["trial"], root)
752
+ if not (suite and trainable_sources(suite_manifest(suite) or {}) and (Path(root) / p["trial"] / "development/rows.json").exists()): continue
753
+ shipped = shipped_temperature(p, root)
754
+ served_t = temperature(p["trial"], root) if shipped is not None else None
755
+ if shipped is not None and abs(served_t - shipped) > SHIPPED_TOLERANCE:
756
+ out.append(f"parent {name} is served at T {served_t:.3f} fitted on {p['trial']}/development ({suite}, a training corpus), "
757
+ f"{abs(served_t - shipped):.3f} from the T {shipped:.3f} its checkpoint ships: its reads are not what it serves (report only)")
758
+ return out
759
+
760
+
761
+ TEMPERATURE_FREE = ("acc",) # the one bootstrapped metric a temperature cannot move (the argmax is invariant)
762
+
763
+
764
+ def _moves_with_temperature(path, rule):
765
+ """Whether a criterion path reads a number the temperature moves: any panel metric but accuracy (a by-length bucket's
766
+ included: acc_16k_plus is accuracy); the unknowable share is scored on records a pool never fits on."""
767
+ panel, metric = (path.split(".") + [""])[:2]
768
+ return panel != "unknowable" and by_length_names(rule["panels"].get(panel, {})).get(metric, metric) not in TEMPERATURE_FREE
769
+
770
+
771
+ def allowlist_problems(tag, suite, sources):
772
+ """A pool read's `sources` allowlist against the sources its suite's manifest lists: a name the suite does not contain
773
+ (a typo) would silently shrink the pool, so it is a problem, as is a suite whose sources cannot be listed."""
774
+ m = suite_manifest(suite_dir(suite)) if suite and suite_dir(suite) else None
775
+ listed = listed_sources(m) if m is not None else set()
776
+ if not listed: return [f"temperature: cannot check the sources allowlist of {tag!r}: {suite!r} lists no sources in this checkout"]
777
+ return [f"temperature: sources for {tag!r} name {unknown} which {suite} does not contain" for unknown in [sorted(set(sources) - listed)] if unknown]
778
+
779
+
780
+ def pool_problems(spec, stages):
781
+ """What is wrong with the round's `temperature` pool: tags that are not reads of the spec, tags an arm cannot locate,
782
+ and pooled tags inside a panel that a temperature-dependent criterion reads (any metric but accuracy, at any stage):
783
+ the temperature would be fitted on the rows it is judged on. The unknowable share is scored on unknowable records,
784
+ which a pool never fits on (knowable rows only), so the unknowable read may be pooled."""
785
+ pool = spec.get("temperature")
786
+ if not pool: return []
787
+ reads, exclude = pool.get("reads") or [], pool.get("exclude_reads", [])
788
+ problems = ["temperature: no reads to pool"] if not reads else []
789
+ ci = pool.get("ci")
790
+ if ci is not None and not (isinstance(ci, dict) and not set(ci) - {"level", "samples", "seed"} and 0 < ci.get("level", 0.9) < 1
791
+ and isinstance(ci.get("samples", SAMPLES), int) and ci.get("samples", SAMPLES) > 0 and isinstance(ci.get("seed", 0), int)):
792
+ problems.append("temperature: ci is {level (0-1), samples (> 0), seed}")
793
+ problems += [f"temperature: {tag!r} is not a read tag of this spec" for tag in reads if tag not in spec["reads"]]
794
+ problems += [f"temperature: exclude_reads {tag!r} is not a read tag of this spec" for tag in exclude if tag not in spec["reads"] and tag != TRANSFER]
795
+ for tag, sources in pool.get("sources", {}).items():
796
+ if tag not in reads: problems.append(f"temperature: sources for {tag!r}, which is not pooled")
797
+ if not sources or not isinstance(sources, list) or not all(isinstance(s, str) for s in sources):
798
+ problems.append(f"temperature: sources for {tag!r} must be a non-empty list of source names"); continue
799
+ problems += allowlist_problems(tag, spec["reads"].get(tag, {}).get("suite"), sources)
800
+ for stage, rule in stages.items():
801
+ for cname, c in rule["criteria"].items():
802
+ for path in _paths(c):
803
+ shared = sorted(set(reads) & set(rule["panels"].get(path.split(".")[0], {}).get("reads", [])))
804
+ if shared and _moves_with_temperature(path, rule):
805
+ problems.append(f"temperature: {shared} pooled, but {'confirm.' + stage if stage else 'rule'} criterion {cname} reads {path}, which the temperature moves")
806
+ for arm, a in spec["arms"].items():
807
+ where = a.get("reads", spec["arm_reads"])
808
+ unplaced = [t for t in [*reads, *exclude] if (t == TRANSFER and not a.get("trial") and transfer_read(spec, arm) is None)
809
+ or (t != TRANSFER and not isinstance(where, str) and t not in where)]
810
+ if unplaced: problems.append(f"arm {arm}: its reads do not locate the temperature pool's {unplaced}")
811
+ return problems
812
+
813
+
814
+ def _validate_plan(study, s, root, partitions):
815
+ from kev.experiment import load_plan, validated_trial
816
+ from kev.suite import read_manifest
817
+ from kev.budget import compute_bound # the admission bound modal_app.admit_study refuses a study over
818
+ problems = []
819
+ try:
820
+ trials = load_plan(root / s["suite"], root / s["plan"]) if partitions else [validated_trial(t, read_manifest(root / s["suite"])) for t in read_json(root / s["plan"])]
821
+ except (ValueError, KeyError) as error:
822
+ return [f"study {study}: plan {s['plan']} does not validate ({type(error).__name__}: {error})"]
823
+ bound = compute_bound(s["gpu"], s["timeout"], len(trials), any(t.get("full_ft") for t in trials))
824
+ if bound > s["budget"]: problems.append(f"study {study}: admission bound ${bound:.2f} exceeds budget ${s['budget']:.2f} (modal_app.admit_study would refuse it)")
825
+ return problems
826
+
827
+
828
+ # --- rule ------------------------------------------------------------------------------------------------------------
829
+
830
+ def compare(candidate, parent, rule, lengths=None):
831
+ """Every panel, the unknowable share and every criterion of one candidate against its parent. Panels whose reads are
832
+ missing on either side are listed under "missing" and their criteria are None; `passed` needs every criterion. A
833
+ candidate served at a pooled temperature records the fit (`temperature_fit`); every candidate records where its
834
+ temperature came from (`temperature_source`); one whose pool rows are missing is reported incomplete with nothing
835
+ computed. lengths(panel) -> {record id: state tokens} serves by_length panels whose rows record no state_tokens."""
836
+ out = {"trial": candidate.trial, "parent": parent.trial, **({"checkpoint": candidate.checkpoint} if candidate.checkpoint else {})}
837
+ if absent := [f"candidate:{d}" for d in candidate.absent_fit()]:
838
+ return {**out, "temperature": None, "missing": absent, "complete": False, "passed": None}
839
+ out.update(temperature=candidate.t, parent_temperature=parent.t, panels={}, missing=[])
840
+ if candidate.fit: out["temperature_fit"] = candidate.fit
841
+ if candidate.pool and candidate.pool.ci: out["temperature_ci"] = pooled_temperature_ci(candidate.pool, candidate.root)
842
+ out["temperature_source"] = candidate.temperature_source()
843
+ out["parent_temperature_source"] = parent.temperature_source() # parents are served at their trial's development rows; `shipped` is head.pt's T
844
+ for name, spec in rule["panels"].items():
845
+ absent = [f"{who}:{side.dirs[t] if t in side.dirs else t}" for who, side in (("candidate", candidate), ("parent", parent)) for t in spec["reads"] if not side.has(t)]
846
+ if isinstance(spec.get("exclude_file"), dict) and not exclude_file_path(spec["exclude_file"], candidate.root).exists():
847
+ absent.append(f"exclude_file:{spec['exclude_file']['path']}")
848
+ if absent: # an optional (report-only) panel's absent reads leave the comparison complete
849
+ if spec.get("optional"): out.setdefault("missing_optional", []).extend(absent)
850
+ else: out["missing"] += absent
851
+ continue
852
+ c, p = candidate.panel(spec), parent.panel(spec)
853
+ mc, mp = metrics(c), metrics(p)
854
+ panel = {"n": len(c)}
855
+ if any(k in spec for k in EXCLUDE_KEYS): # questions the panel's exclusions removed (candidate side; the parent's are the same records)
856
+ panel["excluded"] = len(candidate.panel({k: v for k, v in spec.items() if k not in EXCLUDE_KEYS})) - len(c)
857
+ for m in spec.get("metrics", ()):
858
+ panel[m] = {"candidate": mc[m], "parent": mp[m], **paired(c, p, m)}
859
+ for m in spec.get("report", ()):
860
+ panel.setdefault(m, {"candidate": mc[m], "parent": mp[m]})
861
+ if spec.get("by_length"):
862
+ panel.update(_by_length(spec, c, p, lengths))
863
+ for ref, d in spec.get("versus", {}).items():
864
+ keep = panel_filter(spec, candidate.root)
865
+ rows = [r for r in read_json(candidate.root / d / "rows.json") if r["variant"] == "clean" and keep(r)] # a reference as it served itself
866
+ panel.setdefault("versus", {})[ref] = {m: {"reference": metrics(rows)[m], **paired(c, rows, m)} for m in spec.get("metrics", ())}
867
+ out["panels"][name] = panel
868
+ if rule.get("unknowable"):
869
+ tag = rule["unknowable"]
870
+ if not candidate.has(tag): out["missing"].append(f"candidate:{candidate.dirs[tag]}")
871
+ out["unknowable"] = {"candidate": candidate.unknowable_share(tag), "parent": parent.unknowable_share(tag)}
872
+ out["criteria"] = {name: _criterion(out, c) for name, c in rule["criteria"].items()}
873
+ out["complete"] = not out["missing"]
874
+ out["passed"] = (out["complete"] and all(out["criteria"].values())) if out["criteria"] else None
875
+ return out
876
+
877
+
878
+ def _by_length(spec, candidate, parent, lengths):
879
+ """A by_length panel's entries: "<metric>_<bucket>" -> {candidate, parent, n} (kev.metrics.calibration_by_length of
880
+ each side's served rows, bucketed by the same counts) and "by_length" -> how the tokens were counted."""
881
+ edges, tokenizer = by_length_options(spec)
882
+ counted = all(r.get("state_tokens") is not None for r in (*candidate, *parent))
883
+ if not counted and lengths is None: raise ValueError("a by_length panel needs state-token counts: rows with state_tokens, or lengths")
884
+ counts = None if counted else lengths(spec)
885
+ c, p = calibration_by_length(candidate, counts, edges), calibration_by_length(parent, counts, edges)
886
+ out = {"by_length": {"edges": list(edges), "tokens": "rows" if counted else {"records": "suite", "tokenizer": list(tokenizer)}}}
887
+ for bucket in c:
888
+ for m in LENGTH_METRICS: out[f"{m}_{bucket}"] = {"candidate": c[bucket][m], "parent": p[bucket][m], "n": c[bucket]["n"]}
889
+ return out
890
+
891
+
892
+ def value(report, path):
893
+ """A number from a comparison by path; None if its panel was not read."""
894
+ parts = path.split(".")
895
+ if parts[0] == "unknowable": return (report.get("unknowable") or {}).get(parts[1])
896
+ entry = report["panels"].get(parts[0], {}).get(parts[1])
897
+ if entry is None: return None
898
+ if parts[2] in ("lower", "upper"): return entry["ci95"][0 if parts[2] == "lower" else 1]
899
+ return entry[parts[2]]
900
+
901
+
902
+ def _criterion(report, c):
903
+ left = [value(report, p) for p in (c["left"] if isinstance(c["left"], list) else [c["left"]])]
904
+ right = value(report, c["right"]) if isinstance(c["right"], str) else c["right"]
905
+ if right is None or any(v is None for v in left): return None
906
+ lhs = left[0] - left[1] if len(left) == 2 else left[0]
907
+ return OPS[c["op"]](lhs, right + c["plus"] if "plus" in c else right)
908
+
909
+
910
+ def rank(arms, keys, selectable):
911
+ """Passing selectable arms per size, best first by the rank keys (a key summing several paths adds them left to right)."""
912
+ def score(report):
913
+ out = []
914
+ for key in keys:
915
+ paths = key["by"] if isinstance(key["by"], list) else [key["by"]]
916
+ total = value(report, paths[0])
917
+ for p in paths[1:]: total = total + value(report, p)
918
+ out.append(-total if key.get("order", "desc") == "desc" else total)
919
+ return tuple(out)
920
+ by_size = {}
921
+ for arm, report in arms.items():
922
+ by_size.setdefault(size_of(arm), [])
923
+ if report.get("passed") and arm in selectable: by_size[size_of(arm)].append(arm)
924
+ return {size: sorted(names, key=lambda a: score(arms[a])) for size, names in by_size.items()}
925
+
926
+
927
+ def readout(spec, root=ROOT):
928
+ """The registered rule applied to every arm that has finished training (development rows present); an arm without a
929
+ trial (a checkpoint made from others) is compared once its pool and reads are there (compare lists what is missing)."""
930
+ arms = {}
931
+ for arm, a in spec["arms"].items():
932
+ if a.get("trial") and not (Path(root) / a["trial"] / "development/rows.json").exists():
933
+ arms[arm] = {"trial": a["trial"], "parent": spec["parents"][a["parent"]]["trial"], "missing": [f"candidate:{a['trial']}/development/rows.json"], "complete": False, "passed": None}
934
+ continue
935
+ arms[arm] = compare(arm_side(spec, arm, root), parent_side(spec, arm, root), spec["rule"], lambda panel: panel_lengths(spec, panel))
936
+ ranking = rank(arms, spec["rule"].get("rank", []), {a for a, x in spec["arms"].items() if x.get("select", True)})
937
+ return {"round": spec["round"], "registered": spec["registered"], "drop_ids": sorted(spec.get("drop_ids", [])), "arms": arms,
938
+ "ranking": ranking, "candidates": {size: names[0] if names else None for size, names in ranking.items()}}
939
+
940
+
941
+ def confirm(spec, stage, arm, root=ROOT):
942
+ """One confirmation stage for the chosen arm against its parent, on the stage's reads."""
943
+ out = compare(arm_side(spec, arm, root, stage), parent_side(spec, arm, root, stage), spec["confirm"][stage], lambda panel: panel_lengths(spec, panel))
944
+ return {"round": spec["round"], "stage": stage, "arm": arm, **out}
945
+
946
+
947
+ # --- printing --------------------------------------------------------------------------------------------------------
948
+
949
+ RATES = ("acc", "confident_error_rate", "coverage_at_5pct_error", "coverage_at_1pct_error") # printed in percentage points
950
+
951
+
952
+ def _interval(metric, entry):
953
+ scale, fmt = (100, "+.1f") if metric in RATES else (1, "+.3f")
954
+ return f"{scale * entry['delta']:{fmt}} [{scale * entry['ci95'][0]:{fmt}}, {scale * entry['ci95'][1]:{fmt}}]"
955
+
956
+
957
+ def table(report):
958
+ """One line per comparison: every bootstrapped panel metric (rates in pp), the unknowable share, the verdict, what failed;
959
+ under it, each by_length panel's buckets, and a warning when the arm is served at a temperature fitted on its trial's
960
+ development rows of a training corpus (in distribution: round 19's failure mode)."""
961
+ lines = []
962
+ for arm, r in report.get("arms", {report.get("arm"): report}).items():
963
+ cells = [f"{p}.{m} {_interval(m, e)}" for p, panel in r.get("panels", {}).items() for m, e in panel.items() if isinstance(e, dict) and "ci95" in e]
964
+ if r.get("unknowable"): cells.append(f"unk {r['unknowable']['candidate']}")
965
+ failed = [k for k, v in r.get("criteria", {}).items() if v is False]
966
+ verdict = "incomplete" if not r.get("complete") else {True: "PASS", None: "reported", False: "fail: " + ", ".join(failed)}[r["passed"]]
967
+ lines.append(f"{arm:16} {' | '.join(cells)} -> {verdict}" + (f" (missing {len(r['missing'])})" if r.get("missing") else "")
968
+ + (f" (report-only reads missing {len(r['missing_optional'])})" if r.get("missing_optional") else ""))
969
+ if r.get("temperature_ci"):
970
+ ci = r["temperature_ci"]
971
+ lines.append(f"{'':16} T {r['temperature']:.3f}, {100 * ci['level']:.0f} % bootstrap interval [{ci['lower']:.3f}, {ci['upper']:.3f}] ({ci['questions']} pool questions)")
972
+ for p, panel in r.get("panels", {}).items():
973
+ if "by_length" not in panel: continue
974
+ buckets = [name for name, _, _ in length_buckets(tuple(panel["by_length"]["edges"])) if panel[f"ece_{name}"]["n"]]
975
+ lines.append(f"{'':16} {p} by state length (candidate/parent): " + " | ".join(
976
+ f"{b} n={panel[f'ece_{b}']['n']} acc {panel[f'acc_{b}']['candidate']:.3f}/{panel[f'acc_{b}']['parent']:.3f} ece {panel[f'ece_{b}']['candidate']:.4f}/{panel[f'ece_{b}']['parent']:.4f}" for b in buckets))
977
+ source = r.get("temperature_source") or {}
978
+ if source.get("kind") == "trial development rows" and source.get("training_corpus"):
979
+ lines.append(f"{'':16} !!! {arm} is served at T {r['temperature']:.3f} fitted on {source['rows']}, held-out items of its training corpus {source['suite']}: "
980
+ f"in distribution, {ROUND_19.split(': ', 1)[0]}; not a temperature to ship (register a held-out-datasets `temperature` pool)")
981
+ if "candidates" in report: lines.append(f"candidates {report['candidates']}")
982
+ return "\n".join(lines)
983
+
984
+
985
+ # --- Modal orchestration ---------------------------------------------------------------------------------------------
986
+
987
+ def bench_job(run, suite, name, flags=""):
988
+ """One modal_app.py::benchmarks entry. modal_app.parse_jobs reads the run as everything before the suite, so a pinned
989
+ Hub revision (repo@sha) is safe; the suite and the name must not contain '@' or ',' and flags must start with '--'."""
990
+ if any(c in s for s in (suite, name) for c in "@,") or "," in run or (flags and not flags.startswith("--")):
991
+ raise ValueError(f"cannot encode benchmark job {(run, suite, name, flags)}")
992
+ return "@".join([run, suite, name, *([flags] if flags else [])])
993
+
994
+
995
+ def checkpoint_of(entry):
996
+ """Where a read runs from: the trial's checkpoint on the runs volume (runs/X -> /runs/X/checkpoint) unless declared."""
997
+ return entry["checkpoint"] if "checkpoint" in entry else f"/{entry['trial']}/checkpoint"
998
+
999
+
1000
+ def side_reads(spec, arm, rule, who, side, stage=None):
1001
+ """[(read tag, directory)] a side needs for a stage, by tag: the rule's reads; "transfer" as the arm's transfer_read (a
1002
+ trial's in-trial transfer read is never launched); and, for an arm at the rule stage, its temperature pool."""
1003
+ tags = rule_tags(rule)
1004
+ if who == "candidate" and stage is None and spec.get("temperature"):
1005
+ tags |= {*spec["temperature"]["reads"], *spec["temperature"].get("exclude_reads", [])}
1006
+ reads = {}
1007
+ for tag in tags:
1008
+ if tag != TRANSFER: reads.setdefault(side.dirs[tag], tag)
1009
+ elif who == "candidate" and transfer_read(spec, arm): reads.setdefault(side.dirs[tag], transfer_read(spec, arm))
1010
+ return sorted((read, d) for d, read in reads.items())
1011
+
1012
+
1013
+ def read_commands(spec, arm, stage=None, root=ROOT, sides=("candidate",)):
1014
+ """The `modal run` commands for the missing reads of one arm (and, with sides, its parent): one benchmarks call per side
1015
+ with every suite batched, plus one locked_test call per locked read. Reads that exist locally are skipped, and so are
1016
+ reads of a removed suite (kev.suite.REMOVED_SUITES; the panels they feed stay missing)."""
1017
+ rule = spec["confirm"][stage] if stage else spec["rule"]
1018
+ size, commands = size_of(arm), []
1019
+ for who in sides:
1020
+ side = (arm_side if who == "candidate" else parent_side)(spec, arm, root, stage)
1021
+ run = checkpoint_of(spec["arms"][arm] if who == "candidate" else spec["parents"][spec["arms"][arm]["parent"]])
1022
+ jobs = []
1023
+ for tag, d in side_reads(spec, arm, rule, who, side, stage):
1024
+ if (root / d / "rows.json").exists(): continue
1025
+ r = spec["reads"][tag]
1026
+ if r.get("entrypoint") == "locked_test":
1027
+ if not run.startswith("/runs/"): raise ValueError(f"locked_test reads a trial on the runs volume, not {run}")
1028
+ commands.append(["modal", "run", "modal_app.py::locked_test", "--trial", run.removeprefix("/runs/").removesuffix("/checkpoint"),
1029
+ "--name", Path(d).relative_to("runs/locked").parts[0], "--decision", r["decision"], "--gpu", spec.get("gpu", "H100"),
1030
+ *spec.get("locked_args", {}).get(size, [])])
1031
+ continue
1032
+ if not d.startswith("runs/") or "/" in d.removeprefix("runs/"): raise ValueError(f"{tag}: benchmarks writes runs/<name>, not {d}")
1033
+ if gone := removed_suite(r["suite"]): # never made again; skipped, not raised, so a watcher does not retry it forever
1034
+ print(f"!!! {arm}: not launching read {tag} ({d}): {r['suite']} was removed on {gone['removed']} (kev.suite.REMOVED_SUITES)", flush=True)
1035
+ continue
1036
+ jobs.append(bench_job(run, r["suite"], d.removeprefix("runs/"), r.get("flags", "")))
1037
+ if jobs:
1038
+ timeout = spec.get("read_timeout", {}).get(size)
1039
+ commands.append(["modal", "run", "--detach", "modal_app.py::benchmarks", "--jobs", ",".join(jobs), "--gpu", spec.get("gpu", "H100"),
1040
+ *(["--timeout", str(timeout)] if timeout else [])])
1041
+ return commands
1042
+
1043
+
1044
+ def modal_env(spec):
1045
+ return {**os.environ, **({"KEV_APP_NAME": spec["app"]} if spec.get("app") else {}), **({"KEV_GPU": spec["gpu"]} if spec.get("gpu") else {})}
1046
+
1047
+
1048
+ def launch_commands(commands, spec, log_stem, stagger=STAGGER, run=subprocess.Popen, sleep=time.sleep):
1049
+ """Start each command in the background (log under runs/), `stagger` seconds apart; returns the processes."""
1050
+ procs = []
1051
+ for i, cmd in enumerate(commands):
1052
+ if i: sleep(stagger)
1053
+ log = ROOT / "runs" / f"{log_stem}-{i}.log"
1054
+ print("launch:", " ".join(cmd[:4]), "->", log.relative_to(ROOT), flush=True)
1055
+ with log.open("w", encoding="utf-8") as f:
1056
+ procs.append(run([sys.executable, "-m", *cmd], stdout=f, stderr=subprocess.STDOUT, cwd=ROOT, env=modal_env(spec)))
1057
+ return procs
1058
+
1059
+
1060
+ class InFlight(Exception):
1061
+ """An arm's reads were launched recently and may still be running; they are not launched again yet."""
1062
+
1063
+
1064
+ def launch_arm_reads(spec, arm, stage=None, sides=("candidate",), stagger=STAGGER, launch=None, now=time.time, log=print):
1065
+ """Launch one arm's missing reads once. Under a per-arm lock (runs/.reads-r<N>-<arm>.lock: a watcher and a
1066
+ `launch-reads` cannot both launch the arm), the intent is written first (runs/r<N>-reads-<arm>[-<stage>].json: when and
1067
+ what), then the commands start. Detached benchmarks cannot be seen from here, so while an earlier intent is younger
1068
+ than the reads' timeout the arm counts as in flight: InFlight is raised (the watcher retries on a later pass) instead of
1069
+ a second launch whose job would find /runs/bench/<name> taken. Reads whose rows landed are never relaunched."""
1070
+ tag = f"r{spec['round']}-reads-{arm}" + (f"-{stage}" if stage else "")
1071
+ grace = max(IN_FLIGHT, spec.get("read_timeout", {}).get(size_of(arm), 0))
1072
+ with file_lock(ROOT / "runs" / f".reads-r{spec['round']}-{arm}.lock"):
1073
+ commands = read_commands(spec, arm, stage, ROOT, sides)
1074
+ if not commands: log(f"{arm}: every read has landed"); return []
1075
+ intent = ROOT / "runs" / f"{tag}.json"
1076
+ if intent.exists() and now() - read_json(intent)["launched_at"] < grace:
1077
+ since = read_json(intent)["launched_at"]
1078
+ raise InFlight(f"{arm}: reads launched {int(now() - since)} s ago may still be running; not relaunching for another "
1079
+ f"{int(since + grace - now())} s (if they finished on Modal but were never pulled: modal volume get kev-runs /bench/<name> runs/)")
1080
+ write_json(intent, {"launched_at": now(), "commands": commands}, atomic=True)
1081
+ return (launch or launch_commands)(commands, spec, tag, stagger)
1082
+
1083
+
1084
+ def pull(study, spec):
1085
+ """modal_app.py::pull for one study (modal_app.pull_study holds a per-study lock, so concurrent pulls wait)."""
1086
+ subprocess.run([sys.executable, "-m", "modal", "run", "modal_app.py::pull", "--name", study], check=True, cwd=ROOT, env=modal_env(spec))
1087
+
1088
+
1089
+ def continue_command(spec, study, label):
1090
+ s = spec["studies"][study]
1091
+ return ["modal", "run", "modal_app.py::resume", "--study", study, "--suite", s["suite"], "--transfer", s["transfer"], "--trial", label]
1092
+
1093
+
1094
+ def continue_trial(spec, study, label, run=subprocess.run):
1095
+ """The next attempt of a timed-out trial: modal_app.py::resume --trial (modal_app.continue_full_trial) spawns
1096
+ run_full_trial again, which continues from the trial's last resume point, with the GPU and timeout the study was
1097
+ admitted for, and records the call in runs/<study>.spawn.json; for a pending ledger entry it first adopts the call
1098
+ from the trial's lease (or waits). Returns once the record names a new current call (spawned or adopted). Raises
1099
+ NoContinuation when the study has no attempt ledger, when the ledger has no attempt left (kev.budget.trial_attempts,
1100
+ the count compute_bound admitted; a pending last entry is resolved by the command first) or when the command
1101
+ succeeded without recording a call, and RuntimeError when it refused (a live lease, a young pending entry, the
1102
+ network: the watcher asks again on its next pass). Output in runs/<study>.continue-<label>-<attempt>.log."""
1103
+ from kev.budget import trial_attempts
1104
+ path = ROOT / "runs" / f"{study}.spawn.json"
1105
+ record = read_json(path)
1106
+ if record.get("modal_retries") != 0: raise NoContinuation(f"{path.name} has no attempt ledger (spawned with Modal's retries on, before the ledger)")
1107
+ used, allowed = trial_attempts(record, label)
1108
+ last = record["attempts"][label][-1]
1109
+ if used >= allowed and (last["call"] is not None or last.get("abandoned")):
1110
+ raise NoContinuation(f"{used} of {allowed} attempts used; the admission bound (${record.get('bound_usd')}) counts no more")
1111
+ log = ROOT / "runs" / f"{study}.continue-{label}-{used + 1}.log"
1112
+ print("continue:", " ".join(continue_command(spec, study, label)), "->", log.relative_to(ROOT), flush=True)
1113
+ with log.open("a", encoding="utf-8") as f:
1114
+ done = run([sys.executable, "-m", *continue_command(spec, study, label)], stdout=f, stderr=subprocess.STDOUT, cwd=ROOT, env=modal_env(spec))
1115
+ if read_json(path)["calls"][label] != record["calls"][label]: return
1116
+ if getattr(done, "returncode", 0): raise RuntimeError(f"modal_app.py::resume --trial exited {done.returncode} without a new call (see {log.relative_to(ROOT)})")
1117
+ raise NoContinuation(f"modal_app.py::resume recorded no new call (see {log.relative_to(ROOT)})")
1118
+
1119
+
1120
+ # --- watch -----------------------------------------------------------------------------------------------------------
1121
+
1122
+ NETWORK_MARKERS = ("nodename nor servname", "name or service not known", "temporary failure in name resolution", "unavailable", "connection reset",
1123
+ "connection refused", "network is unreachable", "deadline exceeded")
1124
+
1125
+
1126
+ def transient(error):
1127
+ """Whether polling a call failed because of this machine's network (DNS drop, reset connection, gRPC UNAVAILABLE)
1128
+ rather than because the trial failed. A trial's own exception is re-raised by FunctionCall.get with its original
1129
+ type, so only network types and gRPC transport messages count; anything else is the trial's failure."""
1130
+ try:
1131
+ import modal.exception as me
1132
+ modal_types = (me.ConnectionError, me.ClientClosed)
1133
+ except ImportError:
1134
+ modal_types = ()
1135
+ if isinstance(error, (socket.gaierror, ConnectionError, *modal_types)): return True
1136
+ return type(error).__module__.startswith(("grpclib", "modal")) and any(m in str(error).lower() for m in NETWORK_MARKERS)
1137
+
1138
+
1139
+ class TrialFailed(Exception):
1140
+ """A full-weight trial that failed with an error returns {"failed": ...} (modal_app.failed_trial) instead of raising;
1141
+ poll_modal raises this for it, which watch_studies marks failed like any trial error (never continued)."""
1142
+
1143
+
1144
+ class NoContinuation(Exception):
1145
+ """A timed-out trial that is not continued (no attempt left under its study's admission bound, no attempt ledger, or
1146
+ the continuation spawned no call): watch_studies marks it failed."""
1147
+
1148
+
1149
+ def poll_modal(call_id):
1150
+ """'running' | 'done' | 'timeout' (the call ran out of time: its container was cancelled at the timeout) | 'refused'
1151
+ (a full-weight attempt that found another attempt's fresh lease and did not start: modal_app.TrialLease) for a spawned
1152
+ trial; the trial's exception (or a network error) propagates, and a returned failure is raised as TrialFailed."""
1153
+ import modal
1154
+ try:
1155
+ result = modal.FunctionCall.from_id(call_id).get(timeout=0.5)
1156
+ except TimeoutError: # builtin: no output yet (modal.exception.FunctionTimeoutError is not a builtin TimeoutError)
1157
+ return "running"
1158
+ except modal.exception.FunctionTimeoutError:
1159
+ return "timeout"
1160
+ except modal.exception.OutputExpiredError: # finished long ago; the result is on the volume
1161
+ return "done"
1162
+ if isinstance(result, dict) and "failed" in result: raise TrialFailed(result["failed"])
1163
+ if isinstance(result, dict) and "refused" in result: return "refused"
1164
+ return "done"
1165
+
1166
+
1167
+ class Unmapped(Exception):
1168
+ """A finished call trained no arm of the spec: its reads cannot be launched."""
1169
+
1170
+
1171
+ def watch_studies(studies, on_done, poll=poll_modal, interval=120, max_transient=60, sleep=time.sleep, log=print, root=ROOT, now=time.time, on_timeout=None):
1172
+ """Poll every spawned trial of the studies until each is settled, calling on_done(study, label) for a finished trial
1173
+ until it succeeds; returns the finished calls that map to no arm. State lives in runs/<study>.watch.json, replaced
1174
+ atomically after every change, so a restarted watcher resumes: finished trials are not polled again and launched reads
1175
+ are not launched again. `launching_at` is written before on_done runs; a restart that finds it logs the interrupted
1176
+ launch and calls on_done again, which must check what already happened (kev.rounds.launch_arm_reads does). Network
1177
+ errors while polling are retried (max_transient in a row marks the call failed); an on_done that raises (a pull that
1178
+ lost the network, reads still in flight) is retried on the next pass; one that raises Unmapped leaves the call
1179
+ unlaunched, logged, and settled so the watch can end.
1180
+ A call that ends by a timeout goes to on_timeout(study, label) (watch's: the trial's next attempt, which records its call
1181
+ as the label's current call in runs/<study>.spawn.json); the spawn record is read on every pass, and a label whose
1182
+ call changed (a continuation, from here or by hand) is polled again from its new call. Without on_timeout, or when it
1183
+ raises NoContinuation, the timeout is the trial's failure; any other error from it is retried on the next pass
1184
+ (max_transient in a row marks the trial failed). `continuing_at` is written before on_timeout runs; a restart that
1185
+ finds it and a changed call adopts the call, one that finds the call unchanged asks again (modal_app.continue_full_trial
1186
+ refuses a second attempt while one runs, waits for the old container's lease and counts every attempt in the ledger).
1187
+ A call refused by the trial's lease, and a current attempt with no call recorded (a pending ledger entry), go to
1188
+ on_timeout the same way."""
1189
+ runs = Path(root) / "runs"
1190
+ while True:
1191
+ settled = True
1192
+ calls = {study: read_json(runs / f"{study}.spawn.json")["calls"] for study in studies} # a continuation replaces a label's call
1193
+ for study, study_calls in calls.items():
1194
+ path = runs / f"{study}.watch.json"
1195
+ state = read_json(path) if path.exists() else {"calls": {}}
1196
+ for label, call_id in study_calls.items():
1197
+ s = state["calls"].setdefault(label, {"status": "running", "launched": False, "transient": 0})
1198
+ if s.setdefault("call", call_id) != call_id and s["status"] != "done": # the trial's next attempt
1199
+ log(f"{study}/{label}: polling its continuation {call_id} (was {s['call']})")
1200
+ s.update(call=call_id, status="running", transient=0); s.pop("error", None); s.pop("continuing_at", None)
1201
+ if s["status"] == "running" and call_id is None: # a pending ledger entry: spawned, or about to be, with no call recorded
1202
+ s["status"] = "timeout"; log(f"{study}/{label}: no call recorded for its current attempt (pending); a continuation adopts it from its lease or waits")
1203
+ elif s["status"] == "running":
1204
+ try:
1205
+ s["status"], s["transient"] = poll(call_id), 0
1206
+ except Exception as error: # noqa: BLE001 - a trial's own failure or this machine's network, told apart here
1207
+ if transient(error) and s["transient"] + 1 < max_transient:
1208
+ s["transient"] += 1; log(f"{study}/{label}: network error, retrying ({type(error).__name__}: {str(error)[:120]})")
1209
+ else:
1210
+ s["status"], s["error"] = "failed", f"{type(error).__name__}: {str(error)[:300]}"; log(f"{study}/{label}: FAILED {s['error']}")
1211
+ if s["status"] == "refused": # the attempt found another attempt's fresh lease and never started: it needs a next attempt too
1212
+ s["status"] = "timeout"; log(f"{study}/{label}: its call {call_id} was refused by the trial's lease (another attempt's container was alive)")
1213
+ if s["status"] == "timeout":
1214
+ if on_timeout is None:
1215
+ s["status"], s["error"] = "failed", "FunctionTimeoutError (not continued: no on_timeout)"; log(f"{study}/{label}: FAILED {s['error']}")
1216
+ else:
1217
+ if "continuing_at" in s: log(f"{study}/{label}: a continuation asked for at {s['continuing_at']:.0f} spawned no call it recorded; asking again")
1218
+ s["continuing_at"] = now(); write_json(path, state, atomic=True)
1219
+ try:
1220
+ on_timeout(study, label); log(f"{study}/{label}: timed out; its next attempt was spawned")
1221
+ except NoContinuation as error:
1222
+ s["status"], s["error"] = "failed", f"FunctionTimeoutError, not continued: {error}"; s.pop("continuing_at"); log(f"!!! {study}/{label}: FAILED {s['error']}")
1223
+ except Exception as error: # noqa: BLE001 - the network, or a refusal that may clear; retried on the next pass
1224
+ s["transient"] += 1; s.pop("continuing_at")
1225
+ if s["transient"] >= max_transient:
1226
+ s["status"], s["error"] = "failed", f"FunctionTimeoutError, continuation failed {s['transient']} times: {type(error).__name__}: {str(error)[:300]}"; log(f"!!! {study}/{label}: FAILED {s['error']}")
1227
+ else:
1228
+ log(f"{study}/{label}: continuing failed, retrying next pass ({type(error).__name__}: {str(error)[:300]})")
1229
+ if s["status"] == "done" and not s["launched"] and not s.get("unmapped"):
1230
+ if "launching_at" in s: log(f"{study}/{label}: a launch started at {s['launching_at']:.0f} was interrupted; checking the arm's reads before any relaunch")
1231
+ s["launching_at"] = now(); write_json(path, state, atomic=True)
1232
+ try:
1233
+ on_done(study, label); s["launched"] = True; s.pop("launching_at"); log(f"{study}/{label}: done, reads launched")
1234
+ except Unmapped as error:
1235
+ s["unmapped"] = True; s.pop("launching_at"); log(f"!!! {study}/{label}: finished but {error}; its reads were NOT launched")
1236
+ except Exception as error: # noqa: BLE001 - retried on the next pass; only a dead process leaves launching_at behind
1237
+ s.pop("launching_at"); log(f"{study}/{label}: launching reads failed, retrying next pass ({type(error).__name__}: {str(error)[:300]})")
1238
+ write_json(path, state, atomic=True)
1239
+ settled = settled and (s["status"] == "failed" or s["launched"] or bool(s.get("unmapped")))
1240
+ if settled:
1241
+ return [f"{study}/{label}" for study in calls for label, s in read_json(runs / f"{study}.watch.json")["calls"].items() if s.get("unmapped")]
1242
+ sleep(interval)
1243
+
1244
+
1245
+ def arm_of(spec, study, label):
1246
+ """The arm a spawned call trained (trial directories are <NN>-<label>)."""
1247
+ return next((a for a, x in spec["arms"].items() if x.get("trial", "").startswith(f"runs/{study}/") and Path(x["trial"]).name.split("-", 1)[1] == label), None)
1248
+
1249
+
1250
+ def watch(spec, interval=120, stagger=STAGGER, reads_timeout=6 * 3600):
1251
+ """The round from spawned trials to read-out: when a trial finishes, pull its study and launch that arm's reads
1252
+ (`stagger` apart); once every trial is settled, wait for the reads to land and write the read-out. The wait ends when
1253
+ every read is there, when every benchmarks process this watcher started has exited (a failed read never lands), or
1254
+ after reads_timeout (a restarted watcher has no processes to follow)."""
1255
+ last, procs = [0.0], []
1256
+ def on_done(study, label):
1257
+ arm = arm_of(spec, study, label)
1258
+ if arm is None: raise Unmapped(f"no arm of round {spec['round']} has a trial runs/{study}/<NN>-{label}")
1259
+ pull(study, spec)
1260
+ time.sleep(max(0.0, last[0] + stagger - time.time()))
1261
+ procs.extend(launch_arm_reads(spec, arm, stagger=stagger))
1262
+ last[0] = time.time()
1263
+ unmapped = watch_studies(list(spec.get("studies", {})), on_done, interval=interval, on_timeout=lambda study, label: continue_trial(spec, study, label))
1264
+ deadline = time.time() + reads_timeout
1265
+ finished = [a for a, x in spec["arms"].items() if x.get("trial") and (ROOT / x["trial"] / "result.json").exists()]
1266
+ landed = lambda a: all((ROOT / d / "rows.json").exists() for tag, d in side_reads(spec, a, spec["rule"], "candidate", arm_side(spec, a))
1267
+ if not removed_suite(spec["reads"][tag].get("suite") or "")) # a removed suite's read is never launched
1268
+ while (waiting := [a for a in finished if not landed(a)]) and time.time() < deadline:
1269
+ if procs and all(p.poll() is not None for p in procs): break
1270
+ print(f"waiting for the reads of {waiting}", flush=True); time.sleep(interval)
1271
+ report = write_readout(spec)
1272
+ if unmapped: raise SystemExit(f"finished calls that map to no arm, never read: {unmapped} (fix the spec's arms, then launch-reads)")
1273
+ return report
1274
+
1275
+
1276
+ def launchable(spec, rows=True):
1277
+ """Problems that stop a launch: a recorded round is never relaunched (register a new one), and a new round needs its
1278
+ plans, suites and (for trials) its parents' reads, or its read-out could not be computed."""
1279
+ if spec.get("archive"): return [f"round {spec['round']} is a record (evidence on {spec['archive']}); register a new round to train or read again"]
1280
+ return validate(spec, rows=rows).problems
1281
+
1282
+
1283
+ def study_commands(spec):
1284
+ return {name: ["modal", "run", "modal_app.py::study", "--suite", s["suite"], "--plan", s["plan"], "--name", name, "--transfer", s["transfer"],
1285
+ "--budget", str(s["budget"]), "--timeout", str(s["timeout"]), "--gpu", s["gpu"]] for name, s in spec["studies"].items()}
1286
+
1287
+
1288
+ def launch_studies(spec, stagger=STAGGER, run=subprocess.run, sleep=time.sleep):
1289
+ """modal_app.py::study for every study of the round (each spawns its trials on the deployed app and returns), `stagger`
1290
+ apart, output in runs/<study>.log (a filter on this output can hide the SystemExit that explains a refusal)."""
1291
+ for i, (name, cmd) in enumerate(study_commands(spec).items()):
1292
+ if i: sleep(stagger)
1293
+ print("launch:", " ".join(cmd), flush=True)
1294
+ with (ROOT / "runs" / f"{name}.log").open("w", encoding="utf-8") as f:
1295
+ run([sys.executable, "-m", *cmd], stdout=f, stderr=subprocess.STDOUT, cwd=ROOT, env=modal_env(spec), check=True)
1296
+
1297
+
1298
+ def write_readout(spec, root=ROOT, out=None):
1299
+ report = readout(spec, root)
1300
+ out = Path(out or Path(root) / f"runs/r{spec['round']}-readout"); out.mkdir(parents=True, exist_ok=True)
1301
+ write_json(out / f"round{spec['round']}.json", report); print(table(report))
1302
+ return report
1303
+
1304
+
1305
+ # --- CLI -------------------------------------------------------------------------------------------------------------
1306
+
1307
+ def main(argv=None):
1308
+ ap = argparse.ArgumentParser(prog="python -m kev.rounds", description=__doc__.split("\n\n")[0])
1309
+ sub = ap.add_subparsers(dest="cmd", required=True)
1310
+ for name in ("validate", "launch", "launch-reads", "watch", "readout", "confirm"):
1311
+ p = sub.add_parser(name); p.add_argument("spec")
1312
+ if name in ("validate", "readout", "confirm"): p.add_argument("--root", default=str(ROOT), help="checkout whose runs/ holds the rows (default: this one)")
1313
+ if name == "validate": p.add_argument("--partitions", action="store_true", help="verify the suites' partitions through load_plan (may fetch from the Hub)")
1314
+ if name in ("launch", "launch-reads", "watch"): p.add_argument("--stagger", type=int, default=STAGGER)
1315
+ if name in ("launch", "launch-reads"): p.add_argument("--dry-run", action="store_true", help="print the modal commands only")
1316
+ if name == "launch-reads": p.add_argument("--arms", help="comma-separated arms (default: every arm with a pulled result, and every arm that is a checkpoint without a trial)"); p.add_argument("--parents", action="store_true", help="also the parents' missing reads")
1317
+ if name in ("launch-reads", "confirm"): p.add_argument("--stage", required=name == "confirm"); p.add_argument("--arm", help="the candidate (default: the one the readout names)")
1318
+ if name == "watch": p.add_argument("--interval", type=int, default=120)
1319
+ if name in ("readout", "confirm"): p.add_argument("--out")
1320
+ a = ap.parse_args(argv)
1321
+ spec, root = load(a.spec), Path(getattr(a, "root", ROOT))
1322
+ for line in [*calibration_warnings(spec, root), *continuation_warnings(spec, root)] if a.cmd in ("validate", "launch") else ():
1323
+ print(f"!!! warning: {line}")
1324
+ if a.cmd == "validate":
1325
+ problems, archived = validate(spec, root, partitions=a.partitions)
1326
+ where = f"on {spec['archive']}, not in this checkout" if spec.get("archive") else "not in this checkout"
1327
+ if archived: print(f"archived ({where}):\n " + "\n ".join(archived))
1328
+ print("\n".join(problems) or f"round {spec['round']}: ok ({len(spec['arms'])} arms, {len(spec.get('studies', {}))} studies)")
1329
+ raise SystemExit(1 if problems else 0)
1330
+ if a.cmd == "readout":
1331
+ write_readout(spec, root, a.out); return
1332
+ if a.cmd == "confirm":
1333
+ arm = a.arm or _candidate(spec, root)
1334
+ report = confirm(spec, a.stage, arm, root)
1335
+ out = Path(a.out or root / f"runs/r{spec['round']}-verdict"); out.mkdir(parents=True, exist_ok=True)
1336
+ write_json(out / f"{size_of(arm)}-{a.stage}.json", report); print(table(report)); return
1337
+ problems = launchable(spec, rows=a.cmd != "launch-reads") # launch-reads is how a missing parent read gets made
1338
+ if problems: raise SystemExit("spec does not validate for a launch:\n" + "\n".join(problems))
1339
+ if a.cmd == "launch":
1340
+ if a.dry_run: print("\n".join(" ".join(c) for c in study_commands(spec).values()))
1341
+ else: launch_studies(spec, a.stagger)
1342
+ return
1343
+ if a.cmd == "watch":
1344
+ watch(spec, a.interval, a.stagger); return
1345
+ arms = [a.arm or _candidate(spec, root)] if a.stage else (a.arms.split(",") if a.arms else [x for x, v in spec["arms"].items() if not v.get("trial") or (root / v["trial"] / "result.json").exists()])
1346
+ sides = ("candidate", "parent") if a.parents or a.stage else ("candidate",)
1347
+ if a.dry_run:
1348
+ for arm in arms: print("\n".join(" ".join(c) for c in read_commands(spec, arm, a.stage, root, sides)))
1349
+ return
1350
+ for i, arm in enumerate(arms):
1351
+ if i: time.sleep(a.stagger)
1352
+ try: launch_arm_reads(spec, arm, a.stage, sides, a.stagger)
1353
+ except InFlight as error: print(f"!!! {error}")
1354
+
1355
+
1356
+ def _candidate(spec, root):
1357
+ report = read_json(Path(root) / f"runs/r{spec['round']}-readout/round{spec['round']}.json")
1358
+ chosen = [arm for arm in report["candidates"].values() if arm]
1359
+ if len(chosen) != 1: raise SystemExit(f"the readout names {len(chosen)} candidates ({report['candidates']}); pass --arm")
1360
+ return chosen[0]
1361
+
1362
+
1363
+ if __name__ == "__main__":
1364
+ main()
runtime/kev/serve.py ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FastAPI sidecar for the playground: loads one checkpoint, exposes prefill-only decisions.
2
+
3
+ Run: uv run --extra serve python -m kev.serve --run runs/kev --port 8008
4
+
5
+ TypeSafe-compatible: POST /v1/systemone, GET /v1/models, the `x-typesafe-request-id` response header, and bearer auth
6
+ when KEV_API_KEY is set (unset = open server, the local default). Demo extras: POST /v1/systemone/permute (one Choice
7
+ under several option orders) and POST /v1/systemone/separate (each question in its own pass, for the packed-vs-separate
8
+ comparison). KEV_PREFIX_CACHE / KEV_PREFIX_MIN_TOKENS / KEV_PREFIX_MAX_TOKENS size the state-prefix cache; KEV_DATE_FACTS=1 opts into the
9
+ date preprocessing (api.with_date_facts). A state over kev.model.SERVE_MAX_STATE tokens is refused with a 422 (kev.model.admit);
10
+ KEV_TRUNCATE_STATES=1 reads its first SERVE_MAX_STATE tokens instead, and then every response says whether it did. Backend
11
+ and precision follow LoadOptions (KEV_BACKEND, KEV_DTYPE, ...): on Apple Silicon the hybrid Qwen3.5 checkpoints run on MLX
12
+ by default, elsewhere on torch in bf16.
13
+ """
14
+ import argparse, asyncio, atexit, hmac, os, queue, random, sys, threading, time, traceback, uuid
15
+ from concurrent.futures import Future
16
+ import torch
17
+ from dataclasses import dataclass, field, replace
18
+ from fastapi import FastAPI, HTTPException
19
+ from fastapi.middleware.cors import CORSMiddleware
20
+ from fastapi.responses import JSONResponse
21
+ from pydantic import BaseModel, Field
22
+ from .api import SystemOneRequest, to_record, to_answers, output_tokens, with_date_facts
23
+ from .checkpoint import Checkpoint, LoadOptions, fused_available, is_hub_id
24
+ from .device import default_device, empty_cache, out_of_memory, sync
25
+ from .model import SERVE_MAX_STATE, ContextOverflow, admit
26
+
27
+ PREFIX_CACHE_SIZE = int(os.environ.get("KEV_PREFIX_CACHE", "4")) # states kept (KV + DeltaNet states; attention-only backbones also the state's hidden states); 0 disables
28
+ PREFIX_MIN_TOKENS = os.environ.get("KEV_PREFIX_MIN_TOKENS") # states shorter than this are not cached; default = the model's prefix_min_tokens (0 for hybrid backbones and MLX, 384 for attention-only torch models)
29
+ PREFIX_MAX_TOKENS = int(os.environ.get("KEV_PREFIX_MAX_TOKENS", "65536")) # state tokens the cache holds in all (least recently used evicted first); a longer state is not cached.
30
+ # One 64k state (Kev-27B: ~1.3 GB of keys, values and DeltaNet states), or four 16k ones, not four 64k ones
31
+ DATE_FACTS = os.environ.get("KEV_DATE_FACTS", "0") == "1"
32
+ TRUNCATE_STATES = os.environ.get("KEV_TRUNCATE_STATES", "0") == "1" # unset = a state over SERVE_MAX_STATE tokens gets a 422; 1 = read its first SERVE_MAX_STATE tokens, and every response says whether it did (truncated, usage.state_tokens / state_tokens_used)
33
+ API_KEY = os.environ.get("KEV_API_KEY") # unset = open server; set = require Authorization: Bearer <key>, as the TypeSafe clients always send
34
+ MAX_BATCH = 64 # requests the model thread takes at once (kev.cuda_graphs splits them to fit its buffers)
35
+ MODEL_NAMES = ("kev-latest", "jev-latest") # both names serve this checkpoint; jev-latest is the TypeSafe SDK default model, so an unconfigured client works
36
+
37
+
38
+ @dataclass
39
+ class PrefixCache:
40
+ """State prefixes kept across requests, least recently used first: (state token ids, option_isolation) -> prefix.
41
+ At most `size` states and `max_tokens` state tokens in all; states shorter than min_tokens or longer than max_tokens
42
+ are not cached. A batch keeps (copies) only the new states that will still be here after it, its last distinct ones
43
+ within both bounds: the rest would be evicted by the batch itself."""
44
+ size: int
45
+ min_tokens: int
46
+ max_tokens: int = PREFIX_MAX_TOKENS
47
+ entries: dict = field(default_factory=dict)
48
+ hits: int = 0
49
+ misses: int = 0
50
+ oom_retries: int = 0 # batches that ran out of device memory with states cached, dropped them and ran again (Server._run)
51
+
52
+ def plan(self, encs):
53
+ """-> (key per request, None when its state is not cached; its cached prefix or None; whether to keep a new one)."""
54
+ lengths = [enc["seg"].count(0) for enc in encs]
55
+ keys = [(tuple(enc["ids"][:n]), bool(enc.get("option_isolation"))) if self.size and self.min_tokens <= n <= self.max_tokens else None
56
+ for enc, n in zip(encs, lengths)]
57
+ survivors, tokens = set(), 0
58
+ for key in dict.fromkeys(k for k in reversed(keys) if k is not None): # most recent first, as store() keeps them
59
+ if len(survivors) == self.size or tokens + len(key[0]) > self.max_tokens: break
60
+ survivors.add(key); tokens += len(key[0])
61
+ return keys, [self.entries.get(k) if k is not None else None for k in keys], [k in survivors for k in keys]
62
+
63
+ def over(self, keys):
64
+ """Whether these keys (states) exceed either bound: more than `size` states or `max_tokens` state tokens."""
65
+ return len(keys) > self.size or sum(len(k[0]) for k in keys) > self.max_tokens
66
+
67
+ def make_room(self, keys, cached, keep):
68
+ """Drop now, before the batch runs, the entries this batch's store() will evict anyway, so an old state does not
69
+ stay resident through the passes of the new one that replaces it. store() (re)inserts every hit and every kept new
70
+ state as most recent and evicts the least recently used while over a bound; evicting the same keys first leaves
71
+ the same cache (tests/test_unit.py). A hit dropped here is still held by `cached` for this batch. Kev-4B on MLX:
72
+ a 64k state's cache is 2.2 GB, so a new 64k state after another one peaked with both (runs/mlx-long-states)."""
73
+ order = dict.fromkeys(self.entries) # least recently used first, as store() evicts
74
+ for key, old, k in zip(keys, cached, keep):
75
+ if key is not None and (old is not None or k):
76
+ order.pop(key, None); order[key] = None
77
+ while self.over(order):
78
+ key = next(iter(order)); del order[key]; self.entries.pop(key, None)
79
+
80
+ def store(self, keys, cached, prefixes):
81
+ """Record hits and misses, and (re)insert the batch's prefixes in order: most recently used last."""
82
+ for key, old, new in zip(keys, cached, prefixes):
83
+ if key is None: continue
84
+ self.hits += old is not None; self.misses += old is None
85
+ if new is None: continue
86
+ self.entries.pop(key, None); self.entries[key] = new
87
+ while self.over(self.entries): self.entries.pop(next(iter(self.entries)))
88
+
89
+ def clear(self):
90
+ self.entries.clear()
91
+
92
+
93
+ @dataclass
94
+ class Server:
95
+ """The loaded checkpoint, the state-prefix cache (PrefixCache), and the one model thread that runs every forward pass.
96
+
97
+ Request threads encode their record and queue it; the model thread takes everything queued when it becomes free and
98
+ runs it as one batch (model.probs_batch: with CUDA graphs, shared state and row passes; otherwise one request at a
99
+ time), then answers each request. It also captures pending CUDA graphs when the graphs say so (capture_due). `lock` is
100
+ held around each batch and capture: hold it to use the model directly."""
101
+ checkpoint: Checkpoint
102
+ tok: object
103
+ model: object
104
+ device: str
105
+ lock: threading.Lock = field(default_factory=threading.Lock)
106
+ batches: int = 0
107
+ batched_requests: int = 0
108
+ release_date: str = field(default="") # for the TypeSafe model card; resolved once (may ask the Hub)
109
+ truncate_states: bool = field(default_factory=lambda: TRUNCATE_STATES) # KEV_TRUNCATE_STATES; off: an over-length state is refused (kev.model.admit)
110
+
111
+ def __post_init__(self):
112
+ self.release_date = self.release_date or self.checkpoint.release_date()
113
+ self.prefix_cache = PrefixCache(PREFIX_CACHE_SIZE, int(PREFIX_MIN_TOKENS) if PREFIX_MIN_TOKENS else self.model.prefix_min_tokens)
114
+ self.queue, self.stopping = queue.Queue(), threading.Event()
115
+ # the model thread gives up the GIL at every CUDA sync and waits to get it back while the event loop parses and
116
+ # answers requests; at Python's default 5 ms switch interval those waits stretched a batch's model time ~2x.
117
+ # Process-wide, so close() puts it back.
118
+ self.switch_interval = sys.getswitchinterval()
119
+ sys.setswitchinterval(0.0005)
120
+ self.thread = threading.Thread(target=self._work, name="kev-model", daemon=True)
121
+ self.thread.start()
122
+ atexit.register(self.close) # a daemon thread killed inside a CUDA call at interpreter exit aborts the process
123
+
124
+ def close(self):
125
+ """Stop the model thread after its current batch; requests still queued fail, and so do later ones (submit)."""
126
+ if self.stopping.is_set(): return
127
+ self.stopping.set(); self.thread.join()
128
+ while not self.queue.empty():
129
+ self.queue.get_nowait()[1].set_exception(RuntimeError("the server stopped")); self.queue.task_done()
130
+ sys.setswitchinterval(self.switch_interval)
131
+
132
+ def submit(self, rec):
133
+ """Queue one record for the model thread. -> a Future of (probabilities, stats). The state prefix (tokens up to the
134
+ first question) is cached across requests, so a repeated state only pays for its question rows. latency_ms is the
135
+ model time of the batch the request ran in (not its wait in the queue). A record kev.model.admit refuses (a state
136
+ over SERVE_MAX_STATE tokens unless truncate_states, a question row over SERVE_MAX_BRANCH) is a 422 here, before
137
+ it reaches the model."""
138
+ if self.stopping.is_set(): raise HTTPException(503, "the server is stopping")
139
+ try: enc = admit(self.model, self.tok, rec, truncate=self.truncate_states)
140
+ except ValueError as e: # ContextOverflow, or a record the encoder cannot take
141
+ state = isinstance(e, ContextOverflow) and e.max_state is not None
142
+ raise HTTPException(422, str(e) + (f"; or start the server with KEV_TRUNCATE_STATES=1 to read only its first {e.max_state:,} tokens (responses then say truncated: true)" if state else ""))
143
+ done = Future()
144
+ self.queue.put((enc, done))
145
+ return done
146
+
147
+ def probs(self, rec):
148
+ return self.submit(rec).result()
149
+
150
+ def _work(self):
151
+ graphs = getattr(self.model, "graphs", None)
152
+ while not self.stopping.is_set():
153
+ try: batch = [self.queue.get(timeout=0.05)]
154
+ except queue.Empty:
155
+ if graphs is not None and graphs.capture_due(idle=True):
156
+ with self.lock: graphs.capture_pending(limit=1)
157
+ continue
158
+ while len(batch) < MAX_BATCH:
159
+ try: batch.append(self.queue.get_nowait())
160
+ except queue.Empty: break
161
+ try:
162
+ with self.lock: results = self._run([enc for enc, _ in batch])
163
+ except Exception as e: # every request of the batch gets the error; the thread lives on
164
+ traceback.clear_frames(e.__traceback__) # the batch's futures keep the exception, and its frames held the failed pass's tensors (142 MiB, a 3,578-token state on Kev-0.8B); the traceback keeps its lines
165
+ results = [e] * len(batch)
166
+ for (_, done), result in zip(batch, results):
167
+ (done.set_exception if isinstance(result, Exception) else done.set_result)(result)
168
+ self.queue.task_done()
169
+ if graphs is not None and graphs.capture_due(idle=False):
170
+ with self.lock: graphs.capture_pending(limit=1)
171
+
172
+ def _run(self, encs):
173
+ """One batch through model.probs_batch, with the prefix cache. -> per request (probs, stats). A pass out of device
174
+ memory while states are cached drops the cache and runs once more: the cache only saves time, and kept resident
175
+ it failed every later batch of that size (#75, @vtxyer)."""
176
+ sync(self.device); t = time.time()
177
+ for retry in (False, True):
178
+ keys, cached, keep = self.prefix_cache.plan(encs)
179
+ self.prefix_cache.make_room(keys, cached, keep)
180
+ try: ps, prefixes = self.model.probs_batch(encs, cached, keep); break
181
+ except Exception as e: # states held: the cache's, and this batch's hits make_room already dropped from it
182
+ held = self.prefix_cache.entries or any(c is not None for c in cached)
183
+ if retry or not held or not out_of_memory(e): raise
184
+ cached = None; self.prefix_cache.clear(); self.prefix_cache.oom_retries += 1 # after the except: its traceback holds the failed pass's tensors
185
+ empty_cache(self.device)
186
+ sync(self.device); dt = round((time.time() - t) * 1000, 1)
187
+ self.prefix_cache.store(keys, cached, prefixes)
188
+ self.batches += 1; self.batched_requests += len(encs)
189
+ return [([q.tolist() for q in p], {"tokens": len(enc["ids"]), "state_tokens": enc["state_tokens"], "state_tokens_used": enc["seg"].count(0),
190
+ "latency_ms": dt, "prefix_cache_hit": c is not None}) # state_tokens: the request's state; _used: what the model read (less only when truncated)
191
+ for enc, p, c in zip(encs, ps, cached)]
192
+
193
+ def wait_idle(self):
194
+ """Block until every submitted request is answered and no CUDA graph waits to be captured (benchmarks, warm-up)."""
195
+ self.queue.join()
196
+ graphs = getattr(self.model, "graphs", None)
197
+ while graphs is not None and graphs.capture_due(idle=True): time.sleep(0.01)
198
+ with self.lock: pass # a capture in progress finishes
199
+
200
+ def answer(self, req):
201
+ """The /v1/systemone response body for one request."""
202
+ rec, meta = to_record(prepare(req))
203
+ return self._body(req, meta, *self.probs(rec))
204
+
205
+ async def answer_async(self, req):
206
+ """answer() for the event loop: a request waiting on the model thread holds no worker thread, so a container takes
207
+ as many concurrent requests as its batches can absorb (FastAPI runs sync endpoints on a 40-thread pool)."""
208
+ rec, meta = to_record(prepare(req))
209
+ return self._body(req, meta, *await asyncio.wrap_future(self.submit(rec)))
210
+
211
+ def _body(self, req, meta, ps, m):
212
+ """The TypeSafe body. A server that may truncate (truncate_states) also says, on every response, whether it did:
213
+ `truncated`, and usage.state_tokens (the request's state) / state_tokens_used (what the model read), both counting
214
+ the <state> token. The TypeSafe SDKs ignore fields they do not model, so clients keep parsing."""
215
+ answers = to_answers(ps, meta)
216
+ body = {"model": req.model, "answers": answers, "usage": {"input_tokens": m["tokens"], "output_tokens": output_tokens(self.tok, answers)}, "latency_ms": m["latency_ms"]}
217
+ if self.truncate_states:
218
+ body["usage"].update(state_tokens=m["state_tokens"], state_tokens_used=m["state_tokens_used"])
219
+ body["truncated"] = m["state_tokens"] > m["state_tokens_used"]
220
+ return body
221
+
222
+
223
+ def prepare(req):
224
+ """Opt-in preprocessing applied to every request before the model sees it."""
225
+ return req.model_copy(update={"state": with_date_facts(req.state)}) if DATE_FACTS else req
226
+
227
+
228
+ app = FastAPI(title="kev")
229
+ app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], expose_headers=["x-typesafe-request-id", "server-timing"])
230
+
231
+
232
+ @app.middleware("http")
233
+ async def typesafe(request, call_next):
234
+ """Bearer auth (when API_KEY is set) and the request id every TypeSafe client reads off the response."""
235
+ started = time.perf_counter()
236
+ if API_KEY and request.url.path.startswith("/v1") and not hmac.compare_digest(request.headers.get("authorization", ""), f"Bearer {API_KEY}"):
237
+ resp = JSONResponse({"detail": "missing or invalid API key; send Authorization: Bearer <KEV_API_KEY>"}, 401, {"www-authenticate": "Bearer"})
238
+ else:
239
+ resp = await call_next(request)
240
+ resp.headers["x-typesafe-request-id"] = request.headers.get("x-typesafe-request-id") or uuid.uuid4().hex
241
+ resp.headers["server-timing"] = f"app;dur={(time.perf_counter() - started) * 1000:.1f}" # time inside this process, for telling it from the network
242
+ return resp
243
+
244
+
245
+ def server() -> Server:
246
+ return app.state.server
247
+
248
+
249
+ @app.post("/v1/systemone")
250
+ async def systemone(req: SystemOneRequest):
251
+ """TypeSafe-compatible endpoint: typed questions in, typed answers out, one prefill pass."""
252
+ return await server().answer_async(req)
253
+
254
+
255
+ class PermuteSystemOne(BaseModel):
256
+ request: SystemOneRequest
257
+ question: str
258
+ n_perm: int = Field(default=6, ge=1, le=64) # each order is a forward pass; 0 divided by nothing, unbounded counts ran forever (#30)
259
+ seed: int = 0
260
+
261
+
262
+ @app.post("/v1/systemone/permute")
263
+ def systemone_permute(r: PermuteSystemOne):
264
+ """Re-run one Choice question under n_perm option orders. Returns per-order probabilities keyed by option name."""
265
+ q = r.request.questions.get(r.question)
266
+ if q is None or q.type != "choice": raise HTTPException(422, "question must be an existing choice question")
267
+ rng = random.Random(r.seed); keys = list(q.criteria); runs = []
268
+ for i in range(r.n_perm):
269
+ order = list(keys)
270
+ if i > 0: rng.shuffle(order)
271
+ one = r.request.model_copy(update={"questions": {r.question: q.model_copy(update={"criteria": {k: q.criteria[k] for k in order}})}})
272
+ resp = server().answer(one); a = resp["answers"][r.question]
273
+ runs.append({"order": order, "probabilities": a["probabilities"], "choice": a["choice"], "latency_ms": resp["latency_ms"]})
274
+ spread = {k: max(x["probabilities"][k] for x in runs) - min(x["probabilities"][k] for x in runs) for k in keys}
275
+ return truncation_marks({"runs": runs, "argmax_stable": len({x["choice"] for x in runs}) == 1, "spread": spread}, resp)
276
+
277
+
278
+ @app.post("/v1/systemone/separate")
279
+ def systemone_separate(req: SystemOneRequest):
280
+ """Answer each question in its own request against the same state (N passes). For packed-vs-separate comparison."""
281
+ parts = [server().answer(req.model_copy(update={"questions": {qid: q}})) for qid, q in req.questions.items()]
282
+ answers = {qid: a for p in parts for qid, a in p["answers"].items()}
283
+ body = {"model": req.model, "answers": answers,
284
+ "usage": {"input_tokens": sum(p["usage"]["input_tokens"] for p in parts), "output_tokens": output_tokens(server().tok, answers)},
285
+ "latency_ms": round(sum(p["latency_ms"] for p in parts), 1)}
286
+ return truncation_marks(body, parts[0])
287
+
288
+
289
+ def truncation_marks(body, part):
290
+ """A demo endpoint built from several /v1/systemone answers on one state carries a truncating server's marks
291
+ (Server._body: `truncated`, usage.state_tokens / state_tokens_used) from any one of them: each read the same tokens."""
292
+ if "truncated" in part:
293
+ body["truncated"] = part["truncated"]
294
+ body.setdefault("usage", {}).update({k: part["usage"][k] for k in ("state_tokens", "state_tokens_used")})
295
+ return body
296
+
297
+
298
+ @app.get("/v1/models")
299
+ def models():
300
+ """One TypeSafe model card (name, description, release_date) per accepted model name, plus the Kev serving details
301
+ a client may ignore: the run, the base, the device, the backend and precision, the temperature, prefix-cache stats."""
302
+ s = server()
303
+ ck, meta = s.checkpoint, s.checkpoint.meta
304
+ card = {"description": f"Kev pointer head on {meta.base}, serving {ck.requested} at temperature {s.model.head.temperature:.2f}",
305
+ "release_date": s.release_date,
306
+ "run": ck.requested, "base": meta.base, "lora": meta.lora, "device": s.device, "backend": s.model.backend, "dtype": s.model.dtype,
307
+ "temperature": s.model.head.temperature, "max_state_tokens": SERVE_MAX_STATE, "truncate_states": s.truncate_states,
308
+ "cuda_graphs": graphs.stats() if (graphs := getattr(s.model, "graphs", None)) else None,
309
+ "prefix_cache": {"size": s.prefix_cache.size, "min_state_tokens": s.prefix_cache.min_tokens, "max_tokens": s.prefix_cache.max_tokens, "hits": s.prefix_cache.hits,
310
+ "misses": s.prefix_cache.misses, "cached_states": len(s.prefix_cache.entries), "oom_retries": s.prefix_cache.oom_retries},
311
+ "batches": {"count": s.batches, "requests": s.batched_requests, "queued": s.queue.qsize()}}
312
+ return {"models": [{"name": name, **card} for name in MODEL_NAMES]}
313
+
314
+
315
+ def main():
316
+ ap = argparse.ArgumentParser()
317
+ ap.add_argument("--run", default="runs/kev")
318
+ ap.add_argument("--fallback", default="runs/smoke")
319
+ ap.add_argument("--host", default="127.0.0.1", help="interface to bind; 0.0.0.0 to serve beyond this machine (a container, a VM behind a proxy)")
320
+ ap.add_argument("--port", type=int, default=8008)
321
+ a = ap.parse_args()
322
+ run = a.run if is_hub_id(a.run) or os.path.exists(f"{a.run}/head.pt") else a.fallback
323
+ if run != a.run: print(f"{a.run} not found, falling back to {run}")
324
+ dev = default_device()
325
+ opts = LoadOptions.from_env()
326
+ if dev == "mps" and opts.attn is None: opts = replace(opts, attn="sdpa") # serving default on Apple GPUs (parity measured)
327
+ if dev != "cpu" and opts.dtype is None: opts = replace(opts, dtype=torch.bfloat16) # serving default: 2-4.5x faster than fp32 on an L4, same answers (LoadOptions.dtype); KEV_DTYPE=fp32 for the exact path
328
+ if dev == "cuda" and opts.cuda_graphs is None: opts = replace(opts, cuda_graphs=True) # serving default: a pass is ~2,000 kernel launches, so replaying graphs cuts warm latency several-fold (kev.cuda_graphs); KEV_CUDA_GRAPHS=0 to decline
329
+ fused_default = dev == "cuda" and opts.fused is None
330
+ if fused_default: opts = replace(opts, fused=fused_available()) # serving default: fused Qwen3.5 kernels, ~1/3 less GPU time per batch (kev.fused_qwen35), when fla is installed; KEV_FUSED=0 to decline, KEV_FUSED=1 to insist
331
+ if opts.backend is None: opts = replace(opts, backend="auto") # serving default: MLX for the hybrid Qwen3.5 checkpoints on Apple Silicon (LoadOptions.backend); KEV_BACKEND=torch to decline
332
+ ck = Checkpoint(run)
333
+ tok, model = ck.load(dev, opts)
334
+ if fused_default and not opts.fused and model.hybrid: print("fused Qwen3.5 kernels off: install the flash-linear-attention version kev/fused_qwen35.py pins (FLA_VERSION) to turn them on")
335
+ app.state.server = Server(ck, tok, model, dev)
336
+ print(f"serving {ck.requested} ({ck.path}) on {dev} via {model.backend} ({model.dtype}) {a.host}:{a.port}; states over {SERVE_MAX_STATE:,} tokens "
337
+ f"{'truncated (KEV_TRUNCATE_STATES=1)' if app.state.server.truncate_states else 'refused (422)'}") # /v1/models reports the run as given, not the resolved cache path
338
+ import uvicorn
339
+ uvicorn.run(app, host=a.host, port=a.port)
340
+
341
+
342
+ if __name__ == "__main__":
343
+ main()
runtime/kev/shared_prefix.py ADDED
@@ -0,0 +1,141 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Training forward through a shared state prefix, for the hybrid Qwen3.5 backbones (kev.train --shared_prefix).
2
+
3
+ The row form (kev.model.forward_rows_batch) trains a record with q questions as q causal rows, state + one branch each,
4
+ so the state runs q times. Here each record's state runs once and its branches continue from it, which costs
5
+ state + the branches instead of q x (state + branch). The two are the same computation: a branch token attends to the
6
+ same tokens at the same positions, and it reaches the state through what each layer leaves behind at the end of the
7
+ state, as the layer's own cache path reads it for inference (kev.model._branch_rows_from_prefix): the attention layers'
8
+ keys and values, and the Gated DeltaNet layers' short-convolution window (the conv's left context) and recurrent state
9
+ (the chunked delta rule's `initial_state`, which fla differentiates). Two things differ from serving:
10
+ - `Prefix` holds that per-layer state functionally (transformers' cache layers copy into static buffers in place, which
11
+ autograd cannot see through), so every branch's gradient flows back into the state and accumulates there;
12
+ - each decoder layer runs its state pass and its branch pass inside one step, and gradient checkpointing (when the
13
+ backbone has it on) recomputes the step as a whole, so no cache outlives its layer or is written twice on replay.
14
+ States of a batch are left-padded (the padded positions are zeroed and stay zero: a zero input leaves the DeltaNet state,
15
+ the conv window and the residual stream unchanged), branches right-padded.
16
+ Exactness: fp32 with transformers' reference kernels matches the row form to ~1e-5 (logits and every gradient; tiny
17
+ random Qwen3.5 in tests/test_unit.py, Qwen3.5-0.8B-Base in tests/test_model.py). On CUDA fla's Triton kernels round their
18
+ fp32 dots like TF32 and the shared prefix left-pads states, so it rounds differently from the row form: measured against the
19
+ exact kernels, it is as far from the exact answer as the row form's own batchings are, with no systematic shift.
20
+ """
21
+ import types
22
+
23
+ import torch
24
+ import torch.nn as nn
25
+ import torch.nn.functional as F
26
+ from torch.distributed.fsdp import FSDPModule, register_fsdp_forward_method
27
+ from torch.utils.checkpoint import checkpoint
28
+
29
+
30
+ class Prefix:
31
+ """One decoder layer's view of the shared state, duck-typing the cache calls a Qwen3.5 layer makes. The state pass
32
+ stores what the layer leaves (keys and values, or the conv window and the recurrent state); `branch` gathers it once
33
+ per branch row; the branch pass reads it. Nothing is written in place."""
34
+ record_past = True # keeps a one-token pass off the in-place decode kernel
35
+
36
+ def __init__(self):
37
+ self.keys = self.values = self.conv = None
38
+ self.recurrent_states = {}
39
+
40
+ @property
41
+ def layers(self): # layer code reads `cache.layers[layer_idx].recurrent_states[0]`: every index is this prefix
42
+ return self
43
+
44
+ def __getitem__(self, layer_idx):
45
+ return self
46
+
47
+ def has_previous_state(self, layer_idx=None, state_idx=0):
48
+ return self.conv is not None
49
+
50
+ def update_conv_state(self, conv_input, layer_idx, conv_kernel_size, **kwargs):
51
+ if self.conv is not None: return torch.cat([self.conv, conv_input], -1) # the branch sees the state's last inputs
52
+ self.conv = F.pad(conv_input, (max(0, conv_kernel_size - conv_input.shape[-1]), 0))[..., -conv_kernel_size:]
53
+ return conv_input
54
+
55
+ def update_recurrent_state(self, recurrent_state, layer_idx, **kwargs):
56
+ self.recurrent_states.setdefault(0, recurrent_state) # the state's final state; a branch's own is discarded
57
+ return recurrent_state
58
+
59
+ def update(self, keys, values, layer_idx, *args, **kwargs):
60
+ if self.keys is None:
61
+ self.keys, self.values = keys, values
62
+ return keys, values
63
+ return torch.cat([self.keys, keys], -2), torch.cat([self.values, values], -2)
64
+
65
+ def branch(self, owner):
66
+ """Gather the stored state per branch row (owner[i] = the record of branch i)."""
67
+ pick = lambda t: None if t is None else t.index_select(0, owner)
68
+ self.keys, self.values, self.conv = pick(self.keys), pick(self.values), pick(self.conv)
69
+ self.recurrent_states = {k: pick(v) for k, v in self.recurrent_states.items()}
70
+
71
+
72
+ def _masks(allow, dtype, attn):
73
+ """Additive float mask for eager attention, boolean for SDPA (which also takes a float mask, but only in the query's
74
+ dtype, and autocast changes that)."""
75
+ allow = allow[:, None]
76
+ return allow if attn == "sdpa" else torch.zeros(allow.shape, dtype=dtype, device=allow.device).masked_fill(~allow, torch.finfo(dtype).min)
77
+
78
+
79
+ def branch_hidden(lm, splits, pad_id, device):
80
+ """splits[b] = kev.model.rows_of(record b's encoding). -> per record, per question, the branch's final hidden states
81
+ [branch length, d] in fp32 (what the pointer head reads; the state's own hidden states are never needed)."""
82
+ base = lm.get_base_model() if hasattr(lm, "get_base_model") else lm
83
+ owner = [b for b, (_, _, rows) in enumerate(splits) for _ in rows]
84
+ branches = [r for _, _, rows in splits for r in rows]
85
+ Ls, Lb = max(len(s) for s, _, _ in splits), max(len(r["ids"]) for r in branches)
86
+ s_ids = torch.full((len(splits), Ls), pad_id, device=device); s_pos = torch.zeros_like(s_ids); s_real = torch.zeros_like(s_ids, dtype=torch.bool)
87
+ for b, (ids, pos, _) in enumerate(splits): # left-padded
88
+ s_ids[b, Ls - len(ids):] = torch.tensor(ids, device=device); s_pos[b, Ls - len(ids):] = torch.tensor(pos, device=device); s_real[b, Ls - len(ids):] = True
89
+ b_ids = torch.full((len(branches), Lb), pad_id, device=device); b_pos = torch.zeros_like(b_ids); b_real = torch.zeros_like(b_ids, dtype=torch.bool)
90
+ for i, r in enumerate(branches): # right-padded; their positions already continue the state
91
+ b_ids[i, :len(r["ids"])] = torch.tensor(r["ids"], device=device); b_pos[i, :len(r["pos"])] = torch.tensor(r["pos"], device=device); b_real[i, :len(r["ids"])] = True
92
+ run = _method(base, "kev_shared_prefix", _forward) # on an FSDP2 root: embeddings and final norm gathered as for base(...)
93
+ h_b = run(s_ids, s_pos, s_real, b_ids, b_pos, b_real, torch.tensor(owner, device=device))
94
+ out = [[] for _ in splits]
95
+ for i, b in enumerate(owner): out[b].append(h_b[i, :len(branches[i]["ids"])])
96
+ return out
97
+
98
+
99
+ def _forward(base, s_ids, s_pos, s_real, b_ids, b_pos, b_real, owner):
100
+ """The decoder stack over left-padded states and right-padded branches (owner[i] = the state of branch i)."""
101
+ device, (Ls, Lb) = s_ids.device, (s_ids.shape[1], b_ids.shape[1])
102
+ h_s = base.embed_tokens(s_ids) * s_real[..., None]; h_b = base.embed_tokens(b_ids) * b_real[..., None]
103
+ dtype, attn = h_s.dtype, base.config._attn_implementation
104
+ # attention: a state token sees the real state tokens up to itself; a branch token its record's real state and its
105
+ # own branch so far; a padded query keeps its diagonal so no softmax row is empty (its output is zeroed or unused)
106
+ causal = lambda n: torch.ones(n, n, dtype=torch.bool, device=device).tril()
107
+ eye = lambda n: torch.eye(n, dtype=torch.bool, device=device)
108
+ # States without padding (one per micro-batch, or all of one length) need no state mask under SDPA: None makes it
109
+ # causal (is_causal), the flash kernel's case. A 64k-token state's explicit mask would be 4 GB, kept for the whole
110
+ # backward, and would rule the flash kernel out.
111
+ plain = attn == "sdpa" and bool(s_real.all())
112
+ allow_b = torch.cat([s_real[owner][:, None, :].expand(-1, Lb, -1), (causal(Lb)[None] & b_real[:, None, :]) | eye(Lb)[None]], -1)
113
+ masks_s = {"full_attention": None if plain else _masks((causal(Ls)[None] & s_real[:, None, :]) | eye(Ls)[None], dtype, attn), "linear_attention": s_real.to(dtype)}
114
+ masks_b = {"full_attention": _masks(allow_b, dtype, attn), "linear_attention": b_real.to(dtype)}
115
+ rope_s, rope_b = base.rotary_emb(h_s, s_pos[None].expand(3, -1, -1)), base.rotary_emb(h_b, b_pos[None].expand(3, -1, -1))
116
+
117
+ recompute = base.gradient_checkpointing and base.training and torch.is_grad_enabled()
118
+ for layer in base.layers:
119
+ step = _method(layer, "kev_shared_prefix_step", _step)
120
+ args = (h_s, h_b, (rope_s, rope_b), (masks_s[layer.block_type], masks_b[layer.block_type]), (s_pos, b_pos), s_real, owner)
121
+ h_s, h_b = checkpoint(step, *args, use_reentrant=False) if recompute else step(*args)
122
+ return base.norm(h_b).float()
123
+
124
+
125
+ def _step(layer, h_s, h_b, ropes, masks, positions, s_real, owner):
126
+ """One decoder layer: the state pass, then the branch pass reading what it left (`Prefix`). A single call on the layer
127
+ (both passes call its plain `forward`), so an FSDP2 unit is gathered once for both and checkpointing replays both."""
128
+ prefix = Prefix()
129
+ h_s = layer.forward(h_s, position_embeddings=ropes[0], attention_mask=masks[0], position_ids=positions[0], past_key_values=prefix)
130
+ prefix.branch(owner)
131
+ h_b = layer.forward(h_b, position_embeddings=ropes[1], attention_mask=masks[1], position_ids=positions[1], past_key_values=prefix)
132
+ return h_s * s_real[..., None], h_b
133
+
134
+
135
+ def _method(module, name, fn):
136
+ """`fn` bound to `module` as `name`; on an FSDP2 unit also registered as a forward method, so the unit's parameters
137
+ are gathered around the call and resharded after it, and backward re-gathers them from its outputs' hooks."""
138
+ if not hasattr(module, name):
139
+ setattr(module, name, types.MethodType(fn, module))
140
+ if isinstance(module, FSDPModule): register_fsdp_forward_method(module, name)
141
+ return getattr(module, name)
runtime/kev/study_v3.py ADDED
@@ -0,0 +1,232 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import copy
3
+ import json
4
+ import random
5
+ from collections import Counter, defaultdict
6
+ from pathlib import Path
7
+
8
+ from huggingface_hub import HfApi
9
+
10
+ from kev.composition import DEV_SHAPES, HELD_OUT_KEYS, TEST_SHAPES, TRAIN_SHAPES, canonical, check_group, generate as compose, sample_trees
11
+ from kev.contrastive import ORDINAL_FAMILIES, generate
12
+ from kev.data import materialize
13
+ from kev.model import fits, load_tokenizer
14
+ from kev.suite import SPLITS, digest, load_split, read_manifest, semantic_hash, validate_training, write_json, write_jsonl
15
+
16
+ BASES = ("Qwen/Qwen3-0.6B-Base", "Qwen/Qwen3-4B-Base")
17
+ FAMILIES = ("return_window", "spend_threshold", "age_eligibility", "quantity_limit")
18
+
19
+
20
+ def grouped_split(records, calibration_groups):
21
+ groups = defaultdict(dict)
22
+ for r in records:
23
+ m = r["_meta"]
24
+ groups[m["family"]].setdefault(m["group_id"], []).append(r)
25
+ train, calibration = [], []
26
+ for family in sorted(groups):
27
+ units = list(groups[family].values())
28
+ if len(units) <= calibration_groups:
29
+ raise ValueError("not enough groups per family")
30
+ calibration.extend(r for group in units[:calibration_groups] for r in group)
31
+ train.extend(r for group in units[calibration_groups:] for r in group)
32
+ return train, calibration
33
+
34
+
35
+ def legacy(pairs, seed, families=FAMILIES, source="legacy_policy", excluded=()):
36
+ candidates, _ = generate(3 * pairs, seed, families)
37
+ records, seen, counts = [], set(excluded), Counter()
38
+ for a, b in zip(candidates[::2], candidates[1::2]):
39
+ family = a["_meta"]["family"]
40
+ hashes = {semantic_hash(a), semantic_hash(b)}
41
+ if counts[family] >= pairs or hashes & seen:
42
+ continue
43
+ for r in (a, b):
44
+ m = r["_meta"]
45
+ m.update(source=source, group_id=m["pair_id"], variant="clean")
46
+ m["text_sha256"] = semantic_hash(r)
47
+ records.extend((a, b)); seen.update(hashes); counts[family] += 1
48
+ if any(counts[f] != pairs for f in families):
49
+ raise ValueError("insufficient unique legacy pairs")
50
+ return records
51
+
52
+
53
+ def freeze(out, source="evals/decision-v2", transfer="evals/transfer-v2", public_train=None, synthetic_scale=1, inherit_eval=None,
54
+ random_structures=0, groups_per_structure=8, train_styles=(0, 1), matched_arms=True, legacy_families=FAMILIES, structures_seed=None):
55
+ """random_structures > 0: the compositional arm is generated from that many random rule trees (negation anywhere,
56
+ held-out and locked structures excluded by canonical key) instead of the eight fixed TRAIN_SHAPES."""
57
+ """inherit_eval: a frozen decision suite whose development/test bytes are reused verbatim (only train/calibration are
58
+ regenerated), so results stay comparable across suite versions that differ in training data only."""
59
+ """public_train: optional larger public training pool (a frozen suite dir). Its train+calibration partitions replace
60
+ the source suite's public train/calibration; development/test still come from `source` so results stay comparable."""
61
+ out, source, transfer = Path(out), Path(source), Path(transfer)
62
+ if out.exists():
63
+ raise FileExistsError("v3 destination already exists; choose a new version")
64
+ original = read_manifest(source)
65
+ old_transfer = read_manifest(transfer)
66
+ revisions = {base: HfApi().model_info(base).sha for base in BASES}
67
+ tokenizers = [load_tokenizer(base, revision=sha) for base, sha in revisions.items()]
68
+ parts = {s: [] for s in SPLITS}
69
+ for split in ("train", "calibration", "development"):
70
+ parts[split] = [r for r in load_split(source, split) if r["_meta"]["source"] != "contrastive"]
71
+ if public_train:
72
+ pool = Path(public_train)
73
+ dev_test_states = {r["_meta"]["text_sha256"] for split in ("development", "test") for r in load_split(source, split, allow_test=True)}
74
+ for split in ("train", "calibration"):
75
+ fresh = [r for r in load_split(pool, split) if r["_meta"]["source"] != "contrastive"]
76
+ leaked = [r for r in fresh if r["_meta"]["text_sha256"] in dev_test_states]
77
+ if leaked: raise ValueError(f"{len(leaked)} public {split} records collide with development/test states")
78
+ parts[split] = fresh
79
+ source_dirs = [pool]
80
+ else:
81
+ source_dirs = []
82
+ reserved = set()
83
+ for split in ("train", "calibration", "development"):
84
+ reserved.update(semantic_hash(r) for r in parts[split])
85
+ sources = [p for p in source.glob("*.json*")] + [p for p in transfer.glob("*.json*")] + [p for d in source_dirs for p in d.glob("*.json*")]
86
+ parent_hashes = {str(p): digest(p) for p in sources}
87
+
88
+ def admit_groups(records):
89
+ groups = defaultdict(list)
90
+ for r in records:
91
+ groups[r["_meta"]["group_id"]].append(r)
92
+ for group in groups.values():
93
+ if group[0]["_meta"]["source"] in ("compositional", "composition_holdout"):
94
+ check_group(group)
95
+ hashes = {semantic_hash(r) for r in group}
96
+ if hashes & reserved:
97
+ raise ValueError("semantic state collision across groups; choose a new generation seed")
98
+ for r in group:
99
+ if not fits(materialize(r), *tokenizers):
100
+ raise ValueError("synthetic record exceeds the training context")
101
+ reserved.update(hashes)
102
+ return records
103
+
104
+ # development synthetic records are generated first, against public states only, so they do not depend on the size
105
+ # of the synthetic training arms (dev/test bytes stay identical across suite versions); training groups that would
106
+ # collide with a development state are rejected by admit_groups
107
+ if inherit_eval:
108
+ inherited_dev = load_split(inherit_eval, "development")
109
+ parts["development"] = inherited_dev
110
+ reserved.update(semantic_hash(r) for r in inherited_dev)
111
+ reserved.update(semantic_hash(r) for r in load_split(inherit_eval, "test", allow_test=True))
112
+ dev_synthetic = []
113
+ else:
114
+ dev_synthetic = admit_groups(legacy(12, "v3-control-dev", excluded=reserved)) + admit_groups(compose(4, "v3-composition-dev"))
115
+ old_train, old_cal = grouped_split(admit_groups(legacy(64 * synthetic_scale, "v3-control", families=legacy_families, excluded=reserved)), 8 * synthetic_scale)
116
+ if random_structures:
117
+ trees = {f"rand:{canonical(t)}": t for t in sample_trees(random_structures, f"{structures_seed or out.name}-structures")}
118
+ new_train, new_cal = grouped_split(admit_groups(compose(groups_per_structure, "v7-composition", styles=train_styles, trees=trees)), 1)
119
+ else:
120
+ new_train, new_cal = grouped_split(admit_groups(compose(16 * synthetic_scale, "v3-composition")), 2 * synthetic_scale)
121
+ if matched_arms and (len(old_train) != len(new_train) or len(old_cal) != len(new_cal)):
122
+ raise ValueError("synthetic arms have unequal record budgets")
123
+ parts["train"] += old_train + new_train
124
+ parts["calibration"] += old_cal + new_cal
125
+ parts["development"] += dev_synthetic
126
+ transfer_dev = [r for r in load_split(transfer, "development") if r["_meta"]["source"] != "contrastive"]
127
+ transfer_dev += admit_groups(legacy(20, "v3-legacy-transfer", ("authorization", "deadline"), source="legacy_holdout"))
128
+ transfer_dev += admit_groups(compose(8, "v3-composition-transfer", DEV_SHAPES, styles=(1,), source="composition_holdout"))
129
+ final_extra = admit_groups(compose(8, "v3-locked", TEST_SHAPES, styles=(2,), source="composition_holdout"))
130
+ for split in ("train", "calibration"):
131
+ random.Random(f"v3-{split}").shuffle(parts[split])
132
+ manifest = {"version": 3, "base_revisions": revisions, "dataset_revisions": original["dataset_revisions"],
133
+ "parent_files": parent_hashes, "holdout_sources": [],
134
+ "trainable_sources": sorted({s for s in original["trainable_sources"] if s != "contrastive"}
135
+ | ({s for s in read_manifest(public_train)["trainable_sources"]} if public_train else set())) + ["legacy_policy", "compositional"],
136
+ "eval_only_sources": old_transfer["eval_only_sources"] + ["legacy_holdout", "composition_holdout"],
137
+ "context": original["context"],
138
+ "protocol": {"train_shapes": TRAIN_SHAPES, "transfer_shapes": DEV_SHAPES, "locked_shapes": TEST_SHAPES,
139
+ "train_render_styles": list(train_styles), "locked_render_styles": [2],
140
+ "random_structures": random_structures, "groups_per_structure": groups_per_structure if random_structures else None,
141
+ "structures_seed": f"{structures_seed or out.name}-structures" if random_structures else None,
142
+ "excluded_structure_keys": sorted(HELD_OUT_KEYS), "legacy_families_trainable": list(legacy_families),
143
+ "public_train_records": len(parts["train"]) - len(old_train) - len(new_train),
144
+ "public_train_pool": str(public_train) if public_train else str(source), "synthetic_scale": synthetic_scale,
145
+ "inherited_eval": str(inherit_eval) if inherit_eval else None,
146
+ "synthetic_records_per_arm": len(old_train) if matched_arms else {"legacy_policy": len(old_train), "compositional": len(new_train)},
147
+ "calibration": "shared; stratified by family and group",
148
+ "arm_selection": "train_sources selects public sources plus exactly one synthetic arm",
149
+ "primary": "macro development NLL; transfer and confident-error checks required; no automatic release",
150
+ "legacy_test": "Inherited v2 locked test bytes retained without inspecting examples; only new structures added to transfer test",
151
+ "contamination": "Exact semantic checks among new groups, not fuzzy or pretraining decontamination. Inherited legacy test overlap with new synthetic controls is not certified."},
152
+ "files": {}, "code_hashes": {p.name: digest(p) for p in Path(__file__).parent.glob("*.py")}}
153
+ validate_training(parts["train"], manifest)
154
+ for name, partitions, inherited in (("decision-v3", parts, source),
155
+ ("transfer-v3", {"train": [], "calibration": [], "development": transfer_dev, "test": final_extra}, transfer)):
156
+ folder = out / name
157
+ folder.mkdir(parents=True, exist_ok=False)
158
+ m = copy.deepcopy(manifest)
159
+ if name.startswith("transfer"):
160
+ m.update(trainable_sources=[], holdout_sources=manifest["eval_only_sources"], eval_only=True)
161
+ for split, records in partitions.items():
162
+ path = folder / f"{split}.jsonl"
163
+ payload = b""
164
+ inherited_records = inherited_questions = 0
165
+ if split == "test":
166
+ test_src = Path(inherit_eval) if (inherit_eval and name.startswith("decision")) else inherited
167
+ old = read_manifest(test_src)["files"]["test.jsonl"]
168
+ if digest(test_src / "test.jsonl") != old["sha256"]:
169
+ raise ValueError("inherited test checksum mismatch")
170
+ payload = (test_src / "test.jsonl").read_bytes()
171
+ inherited_records, inherited_questions = old["records"], old["questions"]
172
+ if inherit_eval and name.startswith("decision"): records = [] # the inherited test already contains the locked composition groups
173
+ payload += "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records).encode()
174
+ path.write_bytes(payload)
175
+ m["files"][path.name] = {"sha256": digest(path), "records": inherited_records + len(records),
176
+ "questions": inherited_questions + sum(len(r["questions"]) for r in records)}
177
+ write_json(folder / "manifest.json", m)
178
+ print(name, {s: v["records"] for s, v in m["files"].items()}, flush=True)
179
+ if any(digest(p) != h for p, h in parent_hashes.items()):
180
+ raise ValueError("parent artifacts changed")
181
+ return manifest
182
+
183
+
184
+ def smoke_subset(source, out):
185
+ source, out = Path(source), Path(out)
186
+ manifest = read_manifest(source)
187
+ out.mkdir(parents=True, exist_ok=False)
188
+ manifest["files"] = {}
189
+ for split in SPLITS:
190
+ records = []
191
+ if split != "test":
192
+ groups = defaultdict(list)
193
+ for r in load_split(source, split):
194
+ groups[(r["_meta"]["source"], r["_meta"]["group_id"])].append(r)
195
+ taken = Counter()
196
+ for (name, _), group in groups.items():
197
+ if taken[name] < 2:
198
+ records.extend(group); taken[name] += 1
199
+ path = out / f"{split}.jsonl"
200
+ write_jsonl(path, records)
201
+ manifest["files"][path.name] = {"sha256": digest(path), "records": len(records),
202
+ "questions": sum(len(r["questions"]) for r in records)}
203
+ manifest["smoke_only"] = True
204
+ write_json(out / "manifest.json", manifest)
205
+
206
+
207
+ def main():
208
+ ap = argparse.ArgumentParser()
209
+ ap.add_argument("--out", required=True)
210
+ ap.add_argument("--smoke-from")
211
+ ap.add_argument("--public-train", help="frozen suite whose train/calibration public partitions replace the source's (larger pool)")
212
+ ap.add_argument("--synthetic-scale", type=int, default=1, help="multiply synthetic training pairs/groups per arm (both arms stay equal)")
213
+ ap.add_argument("--inherit-eval", help="decision suite whose development/test bytes are reused verbatim")
214
+ ap.add_argument("--random-structures", type=int, default=0, help="compositional arm from N random rule trees (held-out structures excluded)")
215
+ ap.add_argument("--groups-per-structure", type=int, default=8)
216
+ ap.add_argument("--train-styles", default="0,1", help="rendering styles for training compositional records (2 is locked-test only)")
217
+ ap.add_argument("--unmatched-arms", action="store_true", help="allow the two synthetic arms to differ in size")
218
+ ap.add_argument("--legacy-families", default="v3", choices=["v3", "all"], help="'all' adds the ordinal Score threshold families (deadline stays held out)")
219
+ ap.add_argument("--structures-seed", help="reuse another version's random structures (e.g. v7) so only the legacy arm differs")
220
+ a = ap.parse_args()
221
+ styles = tuple(int(x) for x in a.train_styles.split(","))
222
+ if 2 in styles: ap.error("rendering style 2 is reserved for the locked test")
223
+ if a.smoke_from:
224
+ smoke_subset(a.smoke_from, a.out)
225
+ else:
226
+ freeze(a.out, public_train=a.public_train, synthetic_scale=a.synthetic_scale, inherit_eval=a.inherit_eval,
227
+ random_structures=a.random_structures, groups_per_structure=a.groups_per_structure, train_styles=styles, matched_arms=not a.unmatched_arms,
228
+ legacy_families=FAMILIES + ORDINAL_FAMILIES if a.legacy_families == "all" else FAMILIES, structures_seed=a.structures_seed)
229
+
230
+
231
+ if __name__ == "__main__":
232
+ main()
runtime/kev/suite.py ADDED
@@ -0,0 +1,450 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import copy
3
+ import fcntl
4
+ import hashlib
5
+ import json
6
+ import os
7
+ import random
8
+ import shutil
9
+ from collections import Counter
10
+ from contextlib import contextmanager
11
+ from pathlib import Path
12
+
13
+ from kev.composition import DEV_SHAPES, HELD_OUT_KEYS, TEST_SHAPES, TRAIN_SHAPES
14
+ from kev.contrastive import FAMILIES, generate
15
+ from kev.data import ALL_REPOS, ALL_SOURCES, EVAL_ONLY, REPOS, SOURCES, TRAINABLE, TRANSFER_REPOS, TRANSFER_SOURCES, build, dataset_ref, materialize, source_seed
16
+ from kev.model import (MAX_BRANCH, SERVE_MAX_BRANCH, SERVE_MAX_BRANCH_8K, SERVE_MAX_PACKED, SERVE_MAX_STATE, SERVE_MAX_STATE_8K, fits, load_tokenizer,
17
+ training_context)
18
+
19
+ SPLITS = ("train", "calibration", "development", "test")
20
+ BASES = ("Qwen/Qwen2.5-0.5B", "Qwen/Qwen3-0.6B-Base")
21
+ # the encoder limits every frozen record satisfies, as written into manifests ("context")
22
+ CONTEXT = {**training_context(), "truncate": False}
23
+ # what a manifest records for an eval-only suite frozen as published rather than admitted to the training context
24
+ SERVING_CONTEXT = {"max_state": SERVE_MAX_STATE, "max_branch": SERVE_MAX_BRANCH, "max_packed": SERVE_MAX_PACKED, "truncate": False}
25
+ # the serving context before 64k states: what the suites frozen until then record and were admitted under; their builders
26
+ # (hard-v1, devtools-v1, breadth-v1, long states, semif/typesafe, documents-v1) keep it and kev.model.MAX_TRAIN_STATE_8K,
27
+ # so they still rebuild byte for byte
28
+ SERVING_CONTEXT_8K = {"max_state": SERVE_MAX_STATE_8K, "max_branch": SERVE_MAX_BRANCH_8K, "max_packed": SERVE_MAX_STATE_8K + SERVE_MAX_BRANCH_8K, "truncate": False}
29
+ # Clean records are admitted with this many branch tokens to spare, so the variants that add an option (contrast_cases'
30
+ # none-of-these, training-time none/distractor augmentation) still encode under MAX_BRANCH.
31
+ ADMISSION_BRANCH_HEADROOM = 64
32
+ # Frozen suites are mirrored on the Hub. Manifests (with the sha256 of every partition) and the development/test
33
+ # partitions live in git; large training partitions are fetched from this dataset on first use and verified against
34
+ # the manifest, so the suite hash and every provenance record stay unchanged. A suite whose partitions must never be
35
+ # public (a held-out test set) names its own mirror in the manifest, {"mirror": {"dataset": ..., "revision": ...}},
36
+ # usually the private PRIVATE_DATASET; only its manifest is in git, which publishes the hashes but not the text.
37
+ SUITES_DATASET = "jaredpalmer/kev-suites"
38
+ PRIVATE_DATASET = "jaredpalmer/kev-private-evals"
39
+ SUITES_REVISION = "cc4bac803e73112689ec327ffa481c519cbc7a05"
40
+ # partitions larger than this stay out of git (gitignored; the manifest's sha256 still pins them)
41
+ GIT_LIMIT = 10 * 1024 * 1024
42
+ # the pinned tokenizer suites built for the Qwen3.5 family are admitted and length-counted under (hard-v1, devtools-v1, long states)
43
+ ADMISSION_TOKENIZER = ("Qwen/Qwen3.5-4B-Base", "1001bb4d826a52d1f399e183466143f4da7b741b")
44
+ # programmatic policy sources (kev.study_v3 / kev.contrastive); the trainer's mix ablations treat them as one group
45
+ SYNTHETIC_SOURCES = ("legacy_policy", "compositional", "contrastive")
46
+ # Suites deleted from the repo because they are unsound as a gate: {suite dir: why, when, the last round that read it}.
47
+ # load_split / read_manifest refuse them with the reason; kev.rounds.validate lists a read of one as archived for rounds up
48
+ # to `last_round` (their committed rows under runs/ are the record and still reproduce) and refuses it after that.
49
+ REMOVED_SUITES = {
50
+ "evals/external/scienthoon-v1": {
51
+ "removed": "2026-09-27",
52
+ "last_round": 22,
53
+ "reason": ("unsound as a gate: 291 templated synthetic support tickets x 3 questions; `queue` is saturated (0.948-0.952 "
54
+ "for every 27B), `priority` is unlearnable by construction (its own manifest: the label follows an org rule "
55
+ "absent from the text), and `angry` has 15 of 291 gold labels that contradict the text and turns on ~12 stock "
56
+ "closing phrases with disputed conventions"),
57
+ "record": "PLAN.md (Standing rules; 2026-09-27 note); committed rows under runs/ (e.g. runs/r20-scienthoon)",
58
+ },
59
+ "evals/external/wanli-v2": {
60
+ "removed": "2026-09-30",
61
+ "last_round": 26,
62
+ "reason": ("unsound as a gate: 271 of its 1,002 WANLI test pairs (27 %) are ones WANLI's two crowd annotators labelled "
63
+ "differently, and the published gold is one of the two labels; every Kev scores 49-62 % on those against "
64
+ "64-81 % on the 731 the annotators agreed on, and the 2026-09-27 audit measured split-half r 0.04 across 23 "
65
+ "checkpoints (all within 0.735-0.763), a half-width as wide as its 2 pp bar and ~11 % invalid labels"),
66
+ "record": "PLAN.md (Standing rules; 2026-09-30 note); committed rows under runs/ (e.g. runs/r23-27b-k-w85-wanli2)",
67
+ },
68
+ "evals/external/wanli-v1": {
69
+ "removed": "2026-09-30",
70
+ "last_round": 5,
71
+ "reason": ("the same WANLI test pairs as wanli-v2 (SemIf's 256): 63 of 256 (25 %) are ones WANLI's two crowd annotators "
72
+ "labelled differently, and the published gold is one of the two labels"),
73
+ "record": "PLAN.md (2026-09-30 note); README; committed rows under runs/ (e.g. runs/kev-9b-wanli-v1, runs/jev-wanli-v1)",
74
+ },
75
+ "evals/external/typesafe-v1": {
76
+ "removed": "2026-09-30",
77
+ "last_round": 26,
78
+ "reason": ("unsound as a gate: its gold is not a ground truth but the argmax of the average of two closed frontier models' "
79
+ "answers (evals.typesafe.ai: GPT-6 Astra and Claude Fable 5.1), 13 of its 102 reference distributions put "
80
+ "the answer below 0.75 (the two references split), and on 89 answered questions from 20 cases the split-half "
81
+ "correlation across 20 Kev-27B checkpoints is -0.01 (2026-09-27 audit: -0.27): its differences between "
82
+ "checkpoints are noise"),
83
+ "record": "PLAN.md (Standing rules; 2026-09-30 note); README; committed rows under runs/ (e.g. runs/kev-9b-typesafe-v1)",
84
+ },
85
+ }
86
+
87
+
88
+ class RemovedSuite(ValueError):
89
+ """A read of a suite in REMOVED_SUITES."""
90
+
91
+
92
+ def suite_key(path):
93
+ """'evals/...' for a suite path given relative, absolute or as a container saw it (/root/kev/evals/x); None if none."""
94
+ parts = Path(str(path)).parts
95
+ return str(Path(*parts[parts.index("evals"):])) if "evals" in parts else None
96
+
97
+
98
+ def removed_suite(path):
99
+ """The REMOVED_SUITES entry of a suite path (any form suite_key accepts), or None."""
100
+ key = suite_key(path)
101
+ return REMOVED_SUITES.get(key) if key else None
102
+
103
+
104
+ def refuse_removed(path):
105
+ if entry := removed_suite(path):
106
+ raise RemovedSuite(f"{suite_key(path)} was removed on {entry['removed']}: {entry['reason']}. Rounds up to "
107
+ f"{entry['last_round']} read it; their committed rows are the record ({entry['record']}). Do not read it again.")
108
+
109
+
110
+ def digest(path):
111
+ h = hashlib.sha256()
112
+ with Path(path).open("rb") as stream:
113
+ for block in iter(lambda: stream.read(1024 * 1024), b""):
114
+ h.update(block)
115
+ return h.hexdigest()
116
+
117
+
118
+ def record_digest(record):
119
+ return hashlib.sha256(json.dumps(record, ensure_ascii=False, separators=(",", ":")).encode()).hexdigest()
120
+
121
+
122
+ def normalise_text(text):
123
+ """Casefolded, whitespace runs collapsed to one space: the text two states are compared on for exact deduplication."""
124
+ return " ".join(text.casefold().split())
125
+
126
+
127
+ def text_digest(text):
128
+ """sha256 of normalise_text(text): the `text_sha256` of suite builders (kev.data computes the same inline; it cannot
129
+ import this module). scripts/screen_overlap.py tokenises differently on purpose (words only, for n-gram overlap)."""
130
+ return hashlib.sha256(normalise_text(text).encode()).hexdigest()
131
+
132
+
133
+ # Every JSON/JSONL file this repo writes is UTF-8 with LF line endings, whatever the platform's locale says (issue #12:
134
+ # frozen partitions are sha256-checked byte for byte, and they contain non-ASCII text). Read them the same way.
135
+ ENCODING = "utf-8"
136
+
137
+
138
+ def read_json(path):
139
+ return json.loads(Path(path).read_text(encoding=ENCODING))
140
+
141
+
142
+ def write_json(path, value, atomic=False):
143
+ """atomic: write a sibling temp file and os.replace it over `path`, so a reader (or a restart after a crash mid-write)
144
+ sees the old file or the new one, never a torn one. For state files rewritten in place (kev.rounds' watcher)."""
145
+ text = json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n"
146
+ if not atomic:
147
+ Path(path).write_text(text, encoding=ENCODING); return
148
+ tmp = Path(path).with_name(f".{Path(path).name}.tmp")
149
+ tmp.write_text(text, encoding=ENCODING)
150
+ os.replace(tmp, path)
151
+
152
+
153
+ @contextmanager
154
+ def file_lock(path):
155
+ """Hold an exclusive advisory lock on `path` (created if absent) for the block; a second holder waits. Local
156
+ orchestration only: one pull of a study (modal_app.pull_lock), one launch of an arm's reads (kev.rounds)."""
157
+ Path(path).parent.mkdir(parents=True, exist_ok=True)
158
+ with Path(path).open("a", encoding=ENCODING) as lock:
159
+ fcntl.flock(lock, fcntl.LOCK_EX)
160
+ yield
161
+
162
+
163
+ def read_jsonl(path):
164
+ # split on "\n" only: str.splitlines() also breaks on U+2028, U+2029 and U+0085, which write_jsonl leaves unescaped
165
+ return [json.loads(line) for line in Path(path).read_text(encoding=ENCODING).split("\n") if line.strip()]
166
+
167
+
168
+ def write_jsonl(path, records):
169
+ Path(path).write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records), encoding=ENCODING)
170
+
171
+
172
+ def semantic_hash(r):
173
+ """State hash that ignores sentence order for policy/case records (contrastive, compositional), so the same case
174
+ told in a different order counts as the same state; record_digest of the state otherwise."""
175
+ state = r["state"]
176
+ if isinstance(state, dict) and "policy" in state and "case" in state:
177
+ state = {"policy": state["policy"], "sentences": sorted(s.rstrip(".") for s in state["case"].split(". "))}
178
+ return record_digest(state)
179
+
180
+
181
+ def validate_training(records, manifest):
182
+ """Every training record must come from a source the manifest declares trainable, never from an eval-only or held-out
183
+ one, and compositional records must not use a held-out rule structure."""
184
+ allowed = set(manifest.get("trainable_sources", []))
185
+ forbidden = set(EVAL_ONLY) | set(manifest.get("eval_only_sources", [])) | set(manifest.get("holdout_sources", []))
186
+ for r in records:
187
+ m = r["_meta"]
188
+ if m["source"] in forbidden or (allowed and m["source"] not in allowed):
189
+ raise ValueError(f"eval-only or undeclared training source: {m['source']}")
190
+ if m["source"] == "compositional":
191
+ held_shape = m["family"] in DEV_SHAPES + TEST_SHAPES
192
+ held_structure = m.get("structure") in HELD_OUT_KEYS
193
+ if held_shape or held_structure or (m["family"] not in TRAIN_SHAPES and not m["family"].startswith("rand:")):
194
+ raise ValueError("held-out compositional structure in training")
195
+ if not records:
196
+ raise ValueError("empty training partition")
197
+
198
+
199
+ def read_manifest(directory):
200
+ refuse_removed(directory)
201
+ return read_json(Path(directory) / "manifest.json")
202
+
203
+
204
+ def load_split(directory, split, allow_test=False):
205
+ refuse_removed(directory) # a removed suite is refused with its reason, not a missing-file error
206
+ if split not in SPLITS:
207
+ raise ValueError(f"unknown split: {split}")
208
+ if split == "test" and not allow_test:
209
+ raise ValueError("locked test requires explicit --allow-test; never use it for search")
210
+ directory = Path(directory)
211
+ manifest = read_manifest(directory)
212
+ path = directory / f"{split}.jsonl"
213
+ if not path.exists():
214
+ fetch_partition(directory, path.name)
215
+ if digest(path) != manifest["files"][path.name]["sha256"]:
216
+ raise ValueError(f"suite checksum mismatch: {path}")
217
+ records = read_jsonl(path)
218
+ if len(records) != manifest["files"][path.name]["records"]:
219
+ raise ValueError("suite record count mismatch")
220
+ return records
221
+
222
+
223
+ def fetch_partition(directory, filename):
224
+ """Download one partition of a frozen suite from its Hub mirror into place: the manifest's own "mirror" if it names
225
+ one, else SUITES_DATASET@SUITES_REVISION. The caller verifies the sha256."""
226
+ from huggingface_hub import hf_hub_download
227
+ from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError
228
+ directory = Path(directory).resolve()
229
+ evals_root = next((p for p in directory.parents if p.name == "evals"), None)
230
+ if evals_root is None:
231
+ raise FileNotFoundError(f"{directory / filename} is missing and is not under an evals/ tree")
232
+ relative = directory.relative_to(evals_root) / filename
233
+ mirror = read_manifest(directory).get("mirror")
234
+ repo, revision = (mirror["dataset"], mirror["revision"]) if mirror else (SUITES_DATASET, SUITES_REVISION) # a named mirror pins its own revision
235
+ try:
236
+ cached = hf_hub_download(repo, str(relative), repo_type="dataset", revision=revision)
237
+ except (RepositoryNotFoundError, GatedRepoError) as e: # a private mirror answers "not found" to anyone without access
238
+ raise PermissionError(f"{relative} is only in {repo}, which is missing or private to this account; `hf auth login` "
239
+ "(or HF_TOKEN) with access to it, or ask for it") from e
240
+ shutil.copyfile(cached, directory / filename)
241
+ print(f"fetched {relative} from {repo}@{revision[:10]}", flush=True)
242
+
243
+
244
+ def case_copy(record, variant):
245
+ result = copy.deepcopy(record)
246
+ result["_meta"]["group_id"] = record["_meta"].get("group_id", record["_meta"]["id"]) # bootstrap unit (pair for contrastive)
247
+ result["_meta"]["parent_id"] = record["_meta"]["id"] # the clean record this variant perturbs
248
+ result["_meta"]["id"] += "/" + variant
249
+ result["_meta"]["variant"] = variant
250
+ return result
251
+
252
+
253
+ def contrast_cases(record, seed=0):
254
+ candidates = [(qid, q) for qid, q in record["questions"].items() if q["type"] == "choice" and len(q["criteria"]) >= 3]
255
+ if not candidates:
256
+ return []
257
+ qid, q = candidates[0]
258
+ nk = "none_of_these"
259
+ if nk in q["criteria"]:
260
+ raise ValueError("reserved contrast option collision")
261
+ out = []
262
+ for variant in ("none_present", "none_absent"):
263
+ r = case_copy(record, variant)
264
+ r["questions"] = {qid: copy.deepcopy(q)}
265
+ rq = r["questions"][qid]
266
+ rq["criteria"][nk] = "None of these options describes the answer"
267
+ if variant == "none_absent":
268
+ rq["criteria"].pop(rq["label"])
269
+ rq["label"] = nk
270
+ keys = list(rq["criteria"])
271
+ random.Random(source_seed(seed, record["_meta"]["id"])).shuffle(keys)
272
+ rq["criteria"] = {k: rq["criteria"][k] for k in keys}
273
+ r["_meta"]["none_key"] = nk
274
+ out.append(r)
275
+ r = case_copy(record, "permuted")
276
+ rng = random.Random(source_seed(seed, record["_meta"]["id"]))
277
+ for q in r["questions"].values():
278
+ if q["type"] == "choice":
279
+ keys = list(q["criteria"])
280
+ rng.shuffle(keys)
281
+ q["criteria"] = {k: q["criteria"][k] for k in keys}
282
+ out.append(r)
283
+ return out
284
+
285
+
286
+ def select_unique(records, count, seen, tokenizers, report):
287
+ selected = []
288
+ for record in sorted(records, key=lambda r: r["_meta"]["row_sha256"]):
289
+ report["considered"] += 1
290
+ key = record["_meta"]["text_sha256"]
291
+ if key in seen:
292
+ report["duplicate_state"] += 1
293
+ continue
294
+ rec = materialize(record) # a record that fails request validation is a converter bug: let it abort the freeze
295
+ if not fits(rec, *tokenizers, max_branch=MAX_BRANCH - ADMISSION_BRANCH_HEADROOM):
296
+ report["context_rejected"] += 1
297
+ continue
298
+ record["_meta"].update(group_id=record["_meta"]["id"], variant="clean")
299
+ seen.add(key)
300
+ selected.append(record)
301
+ report["accepted"] += 1
302
+ if len(selected) == count:
303
+ return selected
304
+ raise ValueError(f"only {len(selected)}/{count} records fit the common context policy")
305
+
306
+
307
+ def training_state_hashes(suite_dir):
308
+ """Normalized state hashes of a frozen suite's training + calibration partitions, for exact-match contamination checks."""
309
+ hashes = set()
310
+ for split in ("train", "calibration"):
311
+ for record in load_split(suite_dir, split):
312
+ hashes.add(record["_meta"]["text_sha256"])
313
+ return hashes
314
+
315
+
316
+ def freeze(directory, train=300, calibration=40, development=80, test=80, seed=20260918, holdout=("mnli", "sst5"),
317
+ sources=None, repos=None, exclude_states_from=None, contrastive_pairs=0, contrastive_holdout_families=(), contrastive_eval_pairs=40):
318
+ from huggingface_hub import HfApi
319
+
320
+ sources = SOURCES if sources is None else sources
321
+ repos = REPOS if repos is None else repos
322
+ eval_only = set(holdout) >= set(sources)
323
+ directory = Path(directory)
324
+ if directory.exists():
325
+ raise FileExistsError(f"refusing to overwrite frozen suite {directory}")
326
+ hub = HfApi()
327
+ revisions = {repo: hub.dataset_info(*dataset_ref(repo)[:1], revision=dataset_ref(repo)[1]).sha for repo in set(repos.values())}
328
+ base_revisions = {base: hub.model_info(base).sha for base in BASES}
329
+ tokenizers = [load_tokenizer(base, revision=revision) for base, revision in base_revisions.items()]
330
+ trainable_here = [x for x in sources if x not in holdout]
331
+ if set(trainable_here) & set(EVAL_ONLY):
332
+ raise ValueError(f"eval-only sources cannot be trainable: {sorted(set(trainable_here) & set(EVAL_ONLY))}")
333
+ manifest = {
334
+ "version": 2, "seed": seed, "holdout_sources": list(holdout), "eval_only": eval_only,
335
+ "trainable_sources": trainable_here, "eval_only_sources": [x for x in sources if x in holdout],
336
+ "policy": {"trainable": list(TRAINABLE), "eval_only": list(EVAL_ONLY)},
337
+ "dataset_revisions": revisions, "base_revisions": base_revisions,
338
+ "context": {**CONTEXT, "admission_branch_headroom": ADMISSION_BRANCH_HEADROOM},
339
+ "selection": "Normalized exact-state deduplication across partitions; common tokenizer context admission; no fuzzy decontamination or pretraining-contamination claim.",
340
+ "legacy_checkpoints": "Training/calibration overlap for pre-manifest checkpoints is unknown; exploratory only.",
341
+ "objective": "Negative macro-average clean development NLL, equal weight per task; raw probabilities.",
342
+ "excluded_training_states_from": str(exclude_states_from) if exclude_states_from else None,
343
+ "files": {}, "admission": {},
344
+ }
345
+ partitions = {split: [] for split in SPLITS}
346
+ seen = set(training_state_hashes(exclude_states_from)) if exclude_states_from else set()
347
+ manifest["excluded_training_state_hashes"] = len(seen)
348
+ for source in sources:
349
+ report = Counter()
350
+ test_pool = build(max(3 * (development + test), 600), "test", seed, only=[source], revisions=revisions, sources=sources, repos=repos)
351
+ chosen = select_unique(test_pool, development + test, seen, tokenizers, report)
352
+ partitions["development"].extend(chosen[:development])
353
+ partitions["test"].extend(chosen[development:])
354
+ if source not in holdout:
355
+ train_pool = build(max(3 * (train + calibration), 800), "train", seed, only=[source], revisions=revisions, sources=sources, repos=repos)
356
+ chosen = select_unique(train_pool, train + calibration, seen, tokenizers, report)
357
+ partitions["calibration"].extend(chosen[:calibration])
358
+ partitions["train"].extend(chosen[calibration:])
359
+ manifest["admission"][source] = dict(report)
360
+ print(f"froze {source}: {dict(report)}", flush=True)
361
+ if contrastive_pairs:
362
+ families = list(FAMILIES)
363
+ held = list(contrastive_holdout_families)
364
+ unknown = set(held) - set(families)
365
+ if unknown: raise ValueError(f"unknown contrastive families: {sorted(unknown)}")
366
+ trainable_fams = [f for f in families if f not in held]
367
+ # training/calibration pairs and development/test pairs come from disjoint seeds within trainable families;
368
+ # held-out families appear only in development/test (never trained anywhere)
369
+ train_recs, rep_train = generate(contrastive_pairs, seed=f"{seed}-train", families=trainable_fams) if trainable_fams and not eval_only else ([], {})
370
+ # eval-only suites carry the held-out families only; trainable suites carry their trainable families only
371
+ eval_families = (held or families) if eval_only else trainable_fams
372
+ eval_recs, rep_eval = generate(contrastive_eval_pairs, seed=f"{seed}-eval", families=eval_families)
373
+ for r in train_recs + eval_recs:
374
+ r["_meta"].update(group_id=r["_meta"]["family_id"], variant="clean") # siblings are one bootstrap unit
375
+ if r["_meta"]["text_sha256"] in seen: raise ValueError("contrastive state collides with an existing state")
376
+ seen.update(r["_meta"]["text_sha256"] for r in train_recs + eval_recs)
377
+ n_cal = 2 * max(len(train_recs) // 20, 1) if train_recs else 0
378
+ partitions["calibration"].extend(train_recs[:n_cal]); partitions["train"].extend(train_recs[n_cal:])
379
+ # stratified by family: alternate pairs (siblings adjacent) between development and test
380
+ for i in range(0, len(eval_recs), 2):
381
+ partitions["development" if (i // 2) % 2 == 0 else "test"].extend(eval_recs[i : i + 2])
382
+ manifest["contrastive"] = {"pairs_per_family_train": contrastive_pairs, "pairs_per_family_eval": contrastive_eval_pairs,
383
+ "eval_families": eval_families, "trainable_families": trainable_fams, "eval_only_families": held,
384
+ "train_report": rep_train, "eval_report": rep_eval,
385
+ "note": "labels from rule evaluators; ablation + invariance checks passed for every kept pair; no LLM"}
386
+ if trainable_fams and not eval_only: manifest["trainable_sources"].append("contrastive")
387
+ else: manifest["eval_only_sources"].append("contrastive")
388
+ print(f"contrastive: {len(train_recs)} train/cal records, {len(eval_recs)} dev/test records; held-out families {held}", flush=True)
389
+ for split in ("development", "test"):
390
+ extras = []
391
+ per_source = Counter()
392
+ for record in partitions[split]:
393
+ source = record["_meta"]["source"]
394
+ if per_source[source] < 12:
395
+ variants = contrast_cases(record, seed)
396
+ for variant in variants:
397
+ if not fits(materialize(variant), *tokenizers):
398
+ raise ValueError(f"variant of {record['_meta']['id']} exceeds the training context")
399
+ extras.extend(variants)
400
+ per_source[source] += bool(variants)
401
+ partitions[split].extend(extras)
402
+ directory.mkdir(parents=True)
403
+ for split, records in partitions.items():
404
+ path = directory / f"{split}.jsonl"
405
+ write_jsonl(path, records)
406
+ manifest["files"][path.name] = {"sha256": digest(path), "records": len(records),
407
+ "questions": sum(len(r["questions"]) for r in records)}
408
+ manifest["code_hashes"] = {name: digest(Path(__file__).parent / name) for name in ("data.py", "api.py", "model.py", "suite.py")}
409
+ write_json(directory / "manifest.json", manifest)
410
+ print(json.dumps(manifest["files"], indent=2), flush=True)
411
+
412
+
413
+ def main():
414
+ ap = argparse.ArgumentParser()
415
+ ap.add_argument("--out", required=True)
416
+ ap.add_argument("--train", type=int, default=300)
417
+ ap.add_argument("--calibration", type=int, default=40)
418
+ ap.add_argument("--development", type=int, default=80)
419
+ ap.add_argument("--test", type=int, default=80)
420
+ ap.add_argument("--seed", type=int, default=20260918)
421
+ ap.add_argument("--transfer", action="store_true", help="eval-only suite from the eight transfer-v1 sources (no train/calibration partitions)")
422
+ ap.add_argument("--sources", help="comma-separated source names (any of ALL_SOURCES); trainable ones get train/calibration partitions")
423
+ ap.add_argument("--holdout", help="comma-separated sources kept eval-only within --sources (eval-only policy sources are always held out)")
424
+ ap.add_argument("--exclude-states-from", help="frozen suite whose train+calibration states must not appear here")
425
+ ap.add_argument("--contrastive-pairs", type=int, default=0, help="programmatic contrastive pairs per trainable family for training (0 = none)")
426
+ ap.add_argument("--contrastive-holdout", default="", help="comma-separated contrastive families kept eval-only")
427
+ a = ap.parse_args()
428
+ cfam = tuple(x for x in a.contrastive_holdout.split(",") if x)
429
+ if min(a.development, a.test) < 1:
430
+ ap.error("split sizes must be positive")
431
+ if a.transfer:
432
+ freeze(a.out, 0, 0, a.development, a.test, a.seed, holdout=tuple(TRANSFER_SOURCES), sources=TRANSFER_SOURCES,
433
+ repos=TRANSFER_REPOS, exclude_states_from=a.exclude_states_from, contrastive_pairs=a.contrastive_pairs, contrastive_holdout_families=cfam)
434
+ elif a.sources:
435
+ names = [x for x in a.sources.split(",") if x]
436
+ unknown = set(names) - set(ALL_SOURCES)
437
+ if unknown: ap.error(f"unknown sources: {sorted(unknown)}")
438
+ holdout = tuple(x for x in names if x in EVAL_ONLY or x in (a.holdout or "").split(","))
439
+ if any(x not in holdout for x in names) and min(a.train, a.calibration) < 1:
440
+ ap.error("train and calibration sizes must be positive when a trainable source is included")
441
+ freeze(a.out, a.train, a.calibration, a.development, a.test, a.seed, holdout=holdout,
442
+ sources={k: ALL_SOURCES[k] for k in names}, repos={k: ALL_REPOS[k] for k in names}, exclude_states_from=a.exclude_states_from,
443
+ contrastive_pairs=a.contrastive_pairs, contrastive_holdout_families=cfam)
444
+ else:
445
+ if min(a.train, a.calibration) < 1: ap.error("split sizes must be positive")
446
+ freeze(a.out, a.train, a.calibration, a.development, a.test, a.seed, exclude_states_from=a.exclude_states_from)
447
+
448
+
449
+ if __name__ == "__main__":
450
+ main()