kleeedolinux commited on
Commit
5278d6b
·
0 Parent(s):

Publish Julia 1 model and Python runtime

Browse files
.gitattributes ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz 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/tokenizer.json filter=lfs diff=lfs merge=lfs -text
37
+ julia/router/benchmarks/runtime/*.perf.data filter=lfs diff=lfs merge=lfs -text
38
+ assets/*.png filter=lfs diff=lfs merge=lfs -text
39
+ assets/*.svg filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+
2
+ __pycache__/
3
+ *.pyc
4
+ build/
5
+ *.egg-info/
6
+ benchmark/
7
+ .venv/
8
+ .pytest_cache/
9
+ *.so
README.md ADDED
@@ -0,0 +1,138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ base_model: jhu-clsp/mmBERT-small
4
+ library_name: pytorch
5
+ tags:
6
+ - decision-model
7
+ - text-classification
8
+ - multilingual
9
+ - routing
10
+ ---
11
+
12
+ ![Supersonic Labs logo](assets/supersonic-labs-light.svg)
13
+
14
+ <h1 align="center">Julia 1</h1>
15
+ <p align="center">From context to decisions.</p>
16
+
17
+ ![Julia 1 — monochrome geometric artwork](assets/julia-banner.png)
18
+
19
+ Julia 1 is the **first model in the Julia family and the first released test of our training system**. At 144.3M parameters, it turns a state, a question, and possible answers into one clear decision. The same interface handles classification, routing, ordered scores, and Boolean decisions.
20
+
21
+ ## Evaluation
22
+
23
+ Measured on **2026-09-24** with H200 BF16 inference and strict encoding. These are results for the checkpoint in this repository.
24
+
25
+ | Benchmark | Correct / total | Julia 1 | Jev reference | Difference |
26
+ | --- | ---: | ---: | ---: | ---: |
27
+ | Typed decisions | 1,463 / 2,000 | **73.15%** | 72.70% | +0.45 pp |
28
+ | AG News · 4 labels | 94 / 100 | **94.00%** | 91.00% | +3.00 pp |
29
+ | DAIR Emotion · 6 labels | 86 / 100 | **86.00%** | 48.00% | +38.00 pp |
30
+ | Banking77 pilot · 72 labels | 64 / 100 | 64.00% | **87.00%** | −23.00 pp |
31
+
32
+ Typed accuracy: **Choice 71.33%** (428/600), **Noul 80.67%** (484/600), **Score 68.88%** (551/800).
33
+
34
+ | MASSIVE scenario classification | Correct / total | Accuracy |
35
+ | --- | ---: | ---: |
36
+ | All 52 locales | 110,573 / 154,648 | **71.50%** |
37
+ | Portuguese · pt-PT | 2,565 / 2,974 | **86.25%** |
38
+ | English · en-US | 2,580 / 2,974 | **86.75%** |
39
+
40
+ MASSIVE measures 18 scenario labels across 52 locales with 2,974 examples per locale. It does not measure intent classification or slot filling. The all-locale value is macro accuracy.
41
+
42
+ **Protocol:** [typed-decisions](https://huggingface.co/datasets/LocalLLaMA/typed-decisions) supplies the 400 test cases; classification pilots follow the pinned [Jev benchmark protocol](https://github.com/AbdelStark/jev-benchmarks/tree/0d610cc53e79bcbec691312b0c4adb4a0e371642) using [BTZSC](https://huggingface.co/datasets/btzsc/btzsc). Jev numbers are supplied comparison references, not a new Jev run. Classification pilots cover 100 examples per dataset. Banking uses a 72-label pilot through a ranking/top-16 shortlist, not a native 72-option call. Abstentions count as incorrect.
43
+
44
+ [Full benchmark results](metrics/accuracy-20260924.json) · [Checkpoint and validation provenance](provenance.json).
45
+
46
+ ## Where Julia is accurate — and where it can miss
47
+
48
+ Julia 1 shows what a focused 144.3M-parameter model can do when the task is a **clear choice grounded in supplied context**. It selected the correct answer in 94 of 100 AG News pilot cases and 86 of 100 Emotion pilot cases, and reached 71.50% scenario accuracy across all 52 MASSIVE locales. The typed-decision suite reached 73.15%. These results make Julia useful for testing classification and routing workflows with explicit candidate answers. They also make this first release a concrete test of our training system, with measured strengths and visible gaps rather than a promise of general intelligence.
49
+
50
+ The tradeoff is knowledge and multi-step reasoning. Julia compares the answers you provide; it cannot be counted on to supply missing facts, solve algebraic equations, or carry a long chain of calculations just because the right answer appears among the options. Those abilities are **not established by the evaluations here**. Ambiguous wording, unfamiliar domains, and long label lists can also cause mistakes: the Banking77 pilot reached 64/100 with a shortlist, below its supplied 87/100 reference. The 100-example pilots are encouraging signals, not guarantees for a new workflow. Evaluate the exact questions and options you plan to use.
51
+
52
+ ## Why a decision model for routing?
53
+
54
+ A conventional rule or keyword router needs someone to enumerate phrases and maintain branch order. Julia instead reads the supplied context and the meanings of the candidate answers together. One model can score a new set of labels without adding a new output head for every workflow, and it can handle `choice`, ordered `score`, and Boolean `noul` requests through one API.
55
+
56
+ This is a model-based decision, so its answer depends on the evidence in `state` and the wording of the options. Keep the options clear and distinct, use strict encoding, and evaluate on your own workflow before using it for consequential actions. The hierarchical `Router` can process larger *choice* lists by narrowing groups and reranking survivors, but each native model call still accepts 2–20 options; a grouped result is not a global probability distribution. See [larger choice sets](julia/router/README.md#larger-choice-sets).
57
+
58
+ ## From mmBERT-small to Julia 1
59
+
60
+ Julia 1 starts from [JHU CLSP's mmBERT-small](https://huggingface.co/jhu-clsp/mmBERT-small), a multilingual ModernBERT encoder. The upstream model is a general-purpose encoder released for tasks such as masked-language modeling and downstream fine-tuning. Julia retains that encoder foundation and tokenizer, then adapts the model to score supplied answer options for a typed question about a state. Its checkpoint adds a decision head and is trained on decision-format examples; the released artifact contains the resulting weights and inference code, not the private training pipeline.
61
+
62
+ | | mmBERT-small | Julia 1 |
63
+ | --- | --- | --- |
64
+ | Main interface | Encoder representations / masked-token modeling | `state` + `question` + 2–20 `options` → one selected option |
65
+ | Parameters | About 140M | 144.3M including the decision components |
66
+ | Context | Upstream architecture supports up to 8,192 tokens | Evaluated Julia decision path uses up to 1,024 tokens |
67
+ | Output | Token or encoder features for a downstream task | Scores and a selected answer in the caller's option order |
68
+ | Usage | General multilingual encoder foundation | Specialized finite-choice decisions |
69
+
70
+ The parameter and context figures for mmBERT-small come from its [model card](https://huggingface.co/jhu-clsp/mmBERT-small). Julia's values describe this checkpoint and its evaluated runtime, not a claim that mmBERT-small itself is limited to Julia's interface. Julia is not a chat or text-generation model and is not interchangeable with `AutoModelForMaskedLM`.
71
+
72
+ ## Start here
73
+
74
+ Python 3.11+ is required. CPU inference works with the standard PyTorch installation; no native router build is needed. Download the complete repository, including its 550.5 MiB checkpoint, then install its Python package:
75
+
76
+ ```bash
77
+ python -m pip install huggingface_hub
78
+ python -c "from huggingface_hub import snapshot_download; snapshot_download('SupersonicLabs/Julia-1', local_dir='Julia-1')"
79
+ python -m pip install -e ./Julia-1
80
+ ```
81
+
82
+ Keep the model loaded between requests:
83
+
84
+ ```python
85
+ from julia import load_model
86
+
87
+ engine = load_model(
88
+ "Julia-1",
89
+ device="cpu",
90
+ strict_encoding=True,
91
+ max_length=1024,
92
+ head_length=512,
93
+ )
94
+
95
+ request = {
96
+ "state": "I was charged twice for the same order.",
97
+ "question": "Which team should handle this request?",
98
+ "options": [
99
+ "Billing and payments",
100
+ "Shipping and delivery",
101
+ "Account access",
102
+ ],
103
+ "type": "choice",
104
+ }
105
+
106
+ result = engine.predict([request])[0]
107
+ print(request["options"][result["index"]])
108
+ print(result["probabilities"])
109
+ ```
110
+
111
+ `predict` accepts a list of requests and returns one result per request, in the same order. Each result has an `index` and `probabilities` aligned with the supplied options. Use `engine.predict(requests, probabilities=False)` when only the selected indices are needed. `engine.logits(requests)` returns raw model scores.
112
+
113
+ | Field | Meaning |
114
+ | --- | --- |
115
+ | `state` | Context as text, a JSON object, or a JSON array. |
116
+ | `question` | The decision to make. |
117
+ | `options` | 2–20 nonempty answer descriptions, in the order you want returned. |
118
+ | `type` | `choice` (default), `score` for ordered options, or `noul` for Boolean decisions. |
119
+
120
+ For `noul`, provide exactly two options: **false first, true second**. For `score`, supply options in their intended order. Strict encoding rejects inputs that exceed the model's limits instead of silently truncating them. The evaluated configuration uses a 1,024-token total sequence and a 512-token question-and-options budget; each option may use at most 48 tokens.
121
+
122
+ The returned percentages are **display values, not calibrated certainty**. If the raw top probability exceeds 95% and every other option is below 4.5%, the result displays 100% for the winner. Values below 1% display as 0%, with their mass redistributed proportionally among the remaining options. The selected index still comes from raw logits. For analysis or your own calibration, use `logits()`.
123
+
124
+ CUDA is available with `device="cuda"` when PyTorch sees a BF16-capable GPU. Set `JULIA_CPU_THREADS` before starting Python to adjust CPU threads (default: 4). See [runtime details](julia/router/README.md).
125
+
126
+ ## Limits and deployment notes
127
+
128
+ - The FP32 weights occupy **550.5 MiB**; allow additional memory for the tokenizer and activations. CPU inference needs no GPU.
129
+ - A native request accepts **2–20 options**. Grouped routing can lose the correct answer during narrowing, and its final probabilities cover final candidates only.
130
+ - Benchmark results do not establish accuracy for a new domain, every language, or high-stakes use. The Banking pilot trails its supplied reference.
131
+ - The Python runtime in `julia/` is required. This repository is not a drop-in Transformers text-classification pipeline or a generative model.
132
+ - Download the actual weights, not a Git LFS pointer. Keep checkpoint files unchanged while an engine is loaded.
133
+
134
+ The model artifacts are licensed under **Apache 2.0**. The training pipeline is not included in this repository.
135
+
136
+ ---
137
+
138
+ <p align="center">Supersonic Labs</p>
assets/julia-banner.png ADDED

Git LFS Details

  • SHA256: 6a039468c06177e462c85b944983118e1a5f2a5f3c725b9a7d52659c2af54de0
  • Pointer size: 131 Bytes
  • Size of remote file: 933 kB
assets/supersonic-labs-dark.svg ADDED

Git LFS Details

  • SHA256: ff3acfbd4c5b41cffc1f7c4be56bf49349d2a60ba3166d14e81cd89bdf26c14e
  • Pointer size: 130 Bytes
  • Size of remote file: 16.2 kB
assets/supersonic-labs-light.svg ADDED

Git LFS Details

  • SHA256: 4b17f4411ec052e68d4514afc13f66ea0a5015e6655c9bfa2971156e63e8a16f
  • Pointer size: 130 Bytes
  • Size of remote file: 16.4 kB
encoder/config.json ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "ModernBertForMaskedLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 2,
8
+ "classifier_activation": "gelu",
9
+ "classifier_bias": false,
10
+ "classifier_dropout": 0.0,
11
+ "classifier_pooling": "mean",
12
+ "cls_token_id": 1,
13
+ "decoder_bias": true,
14
+ "deterministic_flash_attn": false,
15
+ "dtype": "float32",
16
+ "embedding_dropout": 0.0,
17
+ "eos_token_id": 1,
18
+ "global_attn_every_n_layers": 3,
19
+ "gradient_checkpointing": false,
20
+ "hidden_activation": "gelu",
21
+ "hidden_size": 384,
22
+ "initializer_cutoff_factor": 2.0,
23
+ "initializer_range": 0.02,
24
+ "intermediate_size": 1152,
25
+ "layer_norm_eps": 1e-05,
26
+ "layer_types": [
27
+ "full_attention",
28
+ "sliding_attention",
29
+ "sliding_attention",
30
+ "full_attention",
31
+ "sliding_attention",
32
+ "sliding_attention",
33
+ "full_attention",
34
+ "sliding_attention",
35
+ "sliding_attention",
36
+ "full_attention",
37
+ "sliding_attention",
38
+ "sliding_attention",
39
+ "full_attention",
40
+ "sliding_attention",
41
+ "sliding_attention",
42
+ "full_attention",
43
+ "sliding_attention",
44
+ "sliding_attention",
45
+ "full_attention",
46
+ "sliding_attention",
47
+ "sliding_attention",
48
+ "full_attention"
49
+ ],
50
+ "local_attention": 128,
51
+ "mask_token_id": 4,
52
+ "max_position_embeddings": 8192,
53
+ "mlp_bias": false,
54
+ "mlp_dropout": 0.0,
55
+ "model_type": "modernbert",
56
+ "norm_bias": false,
57
+ "norm_eps": 1e-05,
58
+ "num_attention_heads": 6,
59
+ "num_hidden_layers": 22,
60
+ "pad_token_id": 0,
61
+ "position_embedding_type": "sans_pos",
62
+ "repad_logits_with_grad": false,
63
+ "rope_parameters": {
64
+ "full_attention": {
65
+ "rope_theta": 160000,
66
+ "rope_type": "default"
67
+ },
68
+ "sliding_attention": {
69
+ "rope_theta": 160000,
70
+ "rope_type": "default"
71
+ }
72
+ },
73
+ "sep_token_id": 1,
74
+ "sparse_pred_ignore_index": -100,
75
+ "sparse_prediction": true,
76
+ "tie_word_embeddings": true,
77
+ "transformers_version": "5.0.0",
78
+ "vocab_size": 256000
79
+ }
inference-policy.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "single_model": true,
3
+ "max_length": 1024,
4
+ "head_length": 512,
5
+ "strict_encoding": true,
6
+ "calibration": null,
7
+ "replacement_qualified": false,
8
+ "step": 500,
9
+ "weights_sha256": "df853bf7fe424420011f3d0c47a05d7341aa9eefa7fb9f203ea4aada4ad95b72"
10
+ }
julia/__init__.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ """Supersonic Labs Julia typed decision models.
2
+
3
+ Use ``load_model`` for a resident CPU or CUDA inference engine.
4
+ """
5
+
6
+ from .inference import load_model
7
+
8
+ __all__ = ['load_model']
julia/cuda.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """CUDA configuration and explicit environment diagnostics for the L40S path."""
2
+ import importlib.metadata
3
+ import json
4
+ import os
5
+ import platform
6
+ import torch
7
+
8
+
9
+ def configure(device_name='cuda', precision='bf16', seed=42):
10
+ device = torch.device(device_name)
11
+ if device.type == 'cuda' and device.index is None:
12
+ device = torch.device('cuda', 0)
13
+ if device.type == 'cuda' and not torch.cuda.is_available():
14
+ raise RuntimeError('CUDA requested but unavailable. Run nvidia-smi and install the CUDA PyTorch wheel; CPU fallback is disabled.')
15
+ if precision == 'bf16' and device.type == 'cuda' and not torch.cuda.is_bf16_supported():
16
+ raise RuntimeError('This CUDA device does not support BF16; select --precision fp32.')
17
+ torch.set_num_threads(int(os.environ.get('JULIA_CPU_THREADS', '4')))
18
+ torch.manual_seed(seed)
19
+ if device.type == 'cuda':
20
+ torch.cuda.set_device(device)
21
+ torch.cuda.manual_seed_all(seed)
22
+ torch.backends.cuda.matmul.allow_tf32 = True
23
+ torch.backends.cudnn.allow_tf32 = True
24
+ torch.cuda.reset_peak_memory_stats(device)
25
+ torch.set_float32_matmul_precision('high')
26
+ return device
27
+
28
+
29
+ def environment(device):
30
+ result = dict(python=platform.python_version(), torch=str(torch.__version__), cuda=torch.version.cuda,
31
+ transformers=importlib.metadata.version('transformers'), device=str(device))
32
+ if device.type == 'cuda':
33
+ prop = torch.cuda.get_device_properties(device)
34
+ result.update(gpu=prop.name, vram_bytes=prop.total_memory,
35
+ compute_capability=list(torch.cuda.get_device_capability(device)))
36
+ return result
37
+
38
+
39
+ def move(batch, device):
40
+ return {k: v.to(device, non_blocking=device.type == 'cuda') for k, v in batch.items()}
41
+
42
+
43
+ def main():
44
+ device = configure()
45
+ x = torch.randn(512, 512, device=device, dtype=torch.bfloat16, requires_grad=True)
46
+ (x @ x.T).float().square().mean().backward()
47
+ torch.cuda.synchronize()
48
+ report = environment(device)
49
+ report['bf16_forward_backward_finite'] = bool(torch.isfinite(x.grad).all())
50
+ print(json.dumps(report, indent=2))
51
+ if not report['bf16_forward_backward_finite']:
52
+ raise RuntimeError('CUDA BF16 diagnostic produced nonfinite gradients')
53
+
54
+
55
+ if __name__ == '__main__':
56
+ main()
julia/data.py ADDED
@@ -0,0 +1,142 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Validated decision JSONL and the upstream-compatible marker serialization."""
2
+ import hashlib
3
+ import json
4
+ import math
5
+ from pathlib import Path
6
+
7
+ import torch
8
+ from torch.utils.data import Dataset
9
+
10
+ QTYPES = {'choice': 0, 'score': 1, 'noul': 2}
11
+
12
+
13
+ def digest(path):
14
+ with Path(path).open('rb') as stream:
15
+ return hashlib.file_digest(stream, 'sha256').hexdigest()
16
+
17
+
18
+ def validate_row(row, line):
19
+ prefix = f'JSONL line {line}: '
20
+ if not isinstance(row, dict):
21
+ raise ValueError(prefix + 'request must be a JSON object')
22
+ if not isinstance(row.get('state'), (str, dict, list)) or not isinstance(row.get('question'), str):
23
+ raise ValueError(prefix + 'state must be text/JSON and question must be text')
24
+ options = row.get('options')
25
+ if not isinstance(options, list) or not 2 <= len(options) <= 20 or not all(isinstance(x, str) and x for x in options):
26
+ raise ValueError(prefix + 'options must contain 2–20 nonempty rendered descriptions')
27
+ if row.get('type', 'choice') not in QTYPES:
28
+ raise ValueError(prefix + 'type must be choice, score, or noul')
29
+ if row.get('type') == 'noul' and len(options) != 2:
30
+ raise ValueError(prefix + 'noul options must be ordered [false, true]')
31
+ if 'target' in row and (type(row['target']) is not int or not 0 <= row['target'] < len(options)):
32
+ raise ValueError(prefix + 'target must index the supplied option list')
33
+ teacher = row.get('teacher_logits')
34
+ if teacher is not None and (
35
+ not isinstance(teacher, list)
36
+ or len(teacher) != len(options)
37
+ or not all(type(x) in (int, float) and math.isfinite(x) for x in teacher)
38
+ ):
39
+ raise ValueError(prefix + 'teacher logits must be finite and match option count/order')
40
+
41
+
42
+ class Decisions(Dataset):
43
+ def __init__(self, path, require_target=True, require_teacher=False):
44
+ self.rows = []
45
+ with open(path) as stream:
46
+ for line, text in enumerate(stream, 1):
47
+ if not text.strip():
48
+ continue
49
+ try:
50
+ row = json.loads(text)
51
+ except json.JSONDecodeError as error:
52
+ raise ValueError(f'{path}:{line}: invalid JSON: {error.msg}') from error
53
+ validate_row(row, line)
54
+ if require_target and 'target' not in row:
55
+ raise ValueError(f'{path}:{line}: training requires target')
56
+ if require_teacher and 'teacher_logits' not in row:
57
+ raise ValueError(f'{path}:{line}: distillation requires teacher_logits')
58
+ self.rows.append(row)
59
+ if not self.rows:
60
+ raise ValueError(f'{path}: dataset contains no rows')
61
+ flags = {'teacher_logits' in row for row in self.rows}
62
+ if len(flags) != 1:
63
+ raise ValueError('Dataset mixes rows with and without teacher logits')
64
+
65
+ def __len__(self):
66
+ return len(self.rows)
67
+
68
+ def __getitem__(self, index):
69
+ return self.rows[index]
70
+
71
+
72
+ def sequence(tokenizer, row, max_length=1024, head_length=256, *, strict=False):
73
+ if head_length + 4 >= max_length:
74
+ raise ValueError('max_length must leave room beyond the question head')
75
+ if any(x is None for x in (tokenizer.mask_token_id, tokenizer.cls_token_id, tokenizer.sep_token_id)):
76
+ raise ValueError('Tokenizer must define MASK, CLS and SEP IDs')
77
+ state = row['state'] if isinstance(row['state'], str) else json.dumps(row['state'], ensure_ascii=False)
78
+ if strict and any(tokenizer.mask_token in text for text in [state, row['question'], *row['options']]):
79
+ raise ValueError('Reserved model marker in request')
80
+ clean = lambda text: text.replace(tokenizer.mask_token, ' ')
81
+ encode = lambda text: tokenizer(text, add_special_tokens=False)['input_ids']
82
+ head = encode(f"{row.get('type', 'choice')} question: {clean(row['question'])}")
83
+ option_ids = [encode(' ' + clean(x)) for x in row['options']]
84
+ if strict and any(len(x) > 48 for x in option_ids):
85
+ raise ValueError('Option exceeds 48-token model contract')
86
+ options = [[tokenizer.mask_token_id] + x[:48] for x in option_ids]
87
+ budget = head_length - sum(map(len, options))
88
+ if budget < 16:
89
+ per_option = max(4, (head_length - 16) // len(options))
90
+ options = [x[:per_option] for x in options]
91
+ budget = head_length - sum(map(len, options))
92
+ if strict and (len(head) > budget or any(len(x) != len(y) + 1 for x, y in zip(options, option_ids))):
93
+ raise ValueError('Question/options exceed lossless head budget')
94
+ ids = [tokenizer.cls_token_id] + head[:max(8, budget)] + [tokenizer.sep_token_id]
95
+ markers = []
96
+ for option in options:
97
+ markers.append(len(ids))
98
+ ids.extend(option)
99
+ ids.append(tokenizer.sep_token_id)
100
+ state_ids = encode(clean(state))
101
+ room = max_length - len(ids) - 1
102
+ if room < 1:
103
+ raise ValueError('Question/options exceed sequence budget; shorten descriptions')
104
+ if strict and len(state_ids) > room:
105
+ raise ValueError('Game state exceeds lossless context budget')
106
+ result = dict(ids=ids + state_ids[:room] + [tokenizer.sep_token_id], markers=markers,
107
+ qtype=QTYPES[row.get('type', 'choice')], truncated=len(state_ids) > room)
108
+ if strict:
109
+ result['option_tokens'] = list(map(len, option_ids))
110
+ return result
111
+
112
+
113
+ class Collator:
114
+ def __init__(self, tokenizer, max_length=512, head_length=256):
115
+ self.tokenizer, self.max_length, self.head_length = tokenizer, max_length, head_length
116
+
117
+ def __call__(self, rows, *, include_targets=True):
118
+ encoded = [row['_encoded'] if '_encoded' in row else sequence(
119
+ self.tokenizer, row, self.max_length, self.head_length) for row in rows]
120
+ length = min(self.max_length, ((max(len(x['ids']) for x in encoded) + 7) // 8) * 8)
121
+ count = max(len(x['markers']) for x in encoded)
122
+ ids = torch.full((len(rows), length), self.tokenizer.pad_token_id, dtype=torch.long)
123
+ attention = torch.zeros_like(ids)
124
+ positions = torch.zeros((len(rows), count), dtype=torch.long)
125
+ mask = torch.zeros_like(positions, dtype=torch.bool)
126
+ has_teacher = include_targets and all('teacher_logits' in row for row in rows)
127
+ teacher = torch.zeros_like(positions, dtype=torch.float32) if has_teacher else None
128
+ for i, (row, item) in enumerate(zip(rows, encoded)):
129
+ n, k = len(item['ids']), len(item['markers'])
130
+ ids[i, :n] = torch.tensor(item['ids'])
131
+ attention[i, :n] = 1
132
+ positions[i, :k] = torch.tensor(item['markers'])
133
+ mask[i, :k] = True
134
+ if has_teacher:
135
+ teacher[i, :k] = torch.tensor(row['teacher_logits'])
136
+ batch = dict(input_ids=ids, attention_mask=attention, marker_pos=positions,
137
+ marker_mask=mask, qtype=torch.tensor([x['qtype'] for x in encoded]))
138
+ if include_targets and all('target' in row for row in rows):
139
+ batch['labels'] = torch.tensor([row['target'] for row in rows])
140
+ if has_teacher:
141
+ batch['teacher_logits'] = teacher
142
+ return batch
julia/inference.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Resident CUDA inference for native Julia and CUDA INT8 checkpoints."""
2
+ import json
3
+ from pathlib import Path
4
+ import torch
5
+ from safetensors.torch import load_file
6
+ from .cuda import configure, move
7
+ from .data import Collator, validate_row
8
+ from .probabilities import display_probabilities
9
+
10
+
11
+ class TransformerEngine:
12
+ def __init__(self, checkpoint, device='cuda', max_length=1024, head_length=256, *, memory_map=True):
13
+ from transformers import AutoModel, AutoTokenizer
14
+ from .model import JuliaDecisionModel
15
+ self.device = configure(device)
16
+ root = Path(checkpoint)
17
+ if (root / 'INCOMPLETE').exists():
18
+ raise ValueError('Refusing to load an incomplete INT8 export')
19
+ self.tokenizer = AutoTokenizer.from_pretrained(root / 'tokenizer', trust_remote_code=False)
20
+ if (root / 'quantization.json').exists():
21
+ if self.device.type != 'cuda':
22
+ raise ValueError('This INT8 checkpoint requires CUDA')
23
+ encoder = AutoModel.from_pretrained(root / 'encoder', device_map={'': self.device.index or 0},
24
+ attn_implementation='sdpa', trust_remote_code=False)
25
+ config = json.loads((root / 'julia_config.json').read_text())
26
+ self.model = JuliaDecisionModel(encoder, **{k: config[k] for k in ('head_layers', 'n_act', 'dropout')})
27
+ heads = load_file(str(root / 'heads.safetensors'))
28
+ expected = {k for k in self.model.state_dict() if not k.startswith('encoder.')}
29
+ if set(heads) != expected:
30
+ raise ValueError('INT8 decision-head tensor keys do not match checkpoint architecture')
31
+ self.model.load_state_dict(heads, strict=False)
32
+ for name, child in self.model.named_children():
33
+ if name != 'encoder':
34
+ child.to(self.device)
35
+ else:
36
+ self.model = JuliaDecisionModel.from_pretrained(root, memory_map=memory_map).to(self.device)
37
+ self.model.eval()
38
+ self.collate = Collator(self.tokenizer, max_length, head_length)
39
+
40
+ @torch.inference_mode()
41
+ def logits(self, rows):
42
+ if not rows:
43
+ return []
44
+ for i, row in enumerate(rows):
45
+ validate_row(row, i + 1)
46
+ batch = move(self.collate(rows, include_targets=False), self.device)
47
+ with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=self.device.type == 'cuda'):
48
+ logits = self.model(**batch)
49
+ logits = logits.cpu()
50
+ if not torch.isfinite(logits).all():
51
+ raise FloatingPointError('Inference returned nonfinite logits')
52
+ return [values[:len(row['options'])].tolist() for values, row in zip(logits, rows)]
53
+
54
+ def predict(self, rows):
55
+ result = []
56
+ for values in self.logits(rows):
57
+ probabilities = torch.tensor(values).softmax(-1)
58
+ result.append(dict(index=int(probabilities.argmax()),
59
+ probabilities=display_probabilities(probabilities.tolist())))
60
+ return result
61
+
62
+
63
+ def load_model(checkpoint, device='cpu', max_length=1024, head_length=256, *, backend=None, **kwargs):
64
+ """Load a Julia checkpoint through the supported resident inference runtime."""
65
+ from .router.engine import FastEngine
66
+ if checkpoint is None:
67
+ raise ValueError('checkpoint is required; pass a local model directory')
68
+ if backend not in (None, 'torch', 'bend', 'bend-dense'):
69
+ raise ValueError('backend must be torch, bend, or bend-dense')
70
+ return FastEngine(checkpoint, device, max_length, head_length,
71
+ transformer_backend=backend, **kwargs)
72
+
73
+
74
+ Engine = load_model # Backward compatible entry point.
julia/model.py ADDED
@@ -0,0 +1,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Marker based decision model and checkpoint serialization."""
2
+ import json
3
+ from pathlib import Path
4
+
5
+ import torch
6
+ from torch import nn
7
+ from torch.nn import functional as F
8
+ from torch.utils.checkpoint import checkpoint
9
+ from safetensors.torch import load_file, save_file
10
+ from transformers import AutoConfig, AutoModel
11
+
12
+
13
+ class JuliaDecisionModel(nn.Module):
14
+ def __init__(self, encoder, head_layers=2, n_act=2, dropout=.1):
15
+ super().__init__()
16
+ self.encoder = encoder
17
+ width = encoder.config.hidden_size
18
+ self.settings = dict(head_layers=head_layers, n_act=n_act, dropout=dropout)
19
+ layer = nn.TransformerEncoderLayer(width, max(1, width // 64), 4 * width,
20
+ dropout, batch_first=True, norm_first=True)
21
+ self.head = nn.TransformerEncoder(layer, head_layers, enable_nested_tensor=False) if head_layers else None
22
+ self.type_emb = nn.Embedding(3, width)
23
+ self.scorer = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, width), nn.GELU(), nn.Linear(width, 1))
24
+ self.act_head = nn.Sequential(nn.Linear(width + 4, 256), nn.GELU(), nn.Linear(256, n_act))
25
+ self.register_buffer('temperature', torch.ones(3))
26
+ self.head_checkpointing = False
27
+ self.marker_only_head = False
28
+ self.encoder.config.reference_compile = False
29
+
30
+ @classmethod
31
+ def from_backbone(cls, path, revision=None, head_layers=2):
32
+ encoder = AutoModel.from_pretrained(path, revision=revision, trust_remote_code=False,
33
+ attn_implementation='sdpa')
34
+ return cls(encoder, head_layers=head_layers)
35
+
36
+ def forward(self, input_ids, attention_mask, marker_pos, marker_mask, qtype, return_actions=False):
37
+ hidden = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
38
+ hidden = hidden + self.type_emb(qtype)[:, None, :]
39
+ # Only option positions (and CLS for the action head) are consumed.
40
+ # The final block can keep full K/V context while avoiding unused queries
41
+ # and feed-forward outputs. Training retains the original dropout path.
42
+ selected = marker_pos
43
+ if return_actions:
44
+ selected = torch.cat((torch.zeros_like(marker_pos[:, :1]), marker_pos), dim=1)
45
+ sparse = (self.marker_only_head and not self.training and self.head is not None
46
+ and len(self.head.layers) > 0
47
+ and type(self.head.layers[-1]) is nn.TransformerEncoderLayer)
48
+ if self.head is not None:
49
+ padding = ~attention_mask.bool()
50
+ for i, layer in enumerate(self.head.layers):
51
+ if sparse and i == len(self.head.layers) - 1:
52
+ hidden = self._selected_head(layer, hidden, selected, attention_mask)
53
+ elif self.head_checkpointing and self.training:
54
+ hidden = checkpoint(layer, hidden, src_key_padding_mask=padding, use_reentrant=False)
55
+ else:
56
+ hidden = layer(hidden, src_key_padding_mask=padding)
57
+ if not sparse:
58
+ positions = selected[:, :, None].expand(-1, -1, hidden.shape[-1])
59
+ hidden = hidden.gather(1, positions)
60
+ markers = hidden[:, 1:] if return_actions else hidden
61
+ scores = self.scorer(markers).squeeze(-1).float()
62
+ scores = scores.masked_fill(~marker_mask, -1e4)
63
+ if not return_actions:
64
+ return scores
65
+ probability = scores.detach().softmax(-1)
66
+ top = probability.topk(min(2, probability.shape[1]), dim=-1).values
67
+ if top.shape[1] == 1:
68
+ top = torch.cat((top, torch.zeros_like(top)), dim=-1)
69
+ count = marker_mask.sum(-1).clamp_min(2).float()
70
+ entropy = -(probability * probability.clamp_min(1e-9).log()).sum(-1) / count.log()
71
+ features = torch.stack((top[:, 0], top[:, 0] - top[:, 1], entropy, count / 255), -1)
72
+ actions = self.act_head(torch.cat((hidden[:, 0].float(), features), -1))
73
+ return scores, actions
74
+
75
+ @staticmethod
76
+ def _selected_head(layer, hidden, selected, attention_mask):
77
+ """Exact pre-norm block restricted to selected output positions (eval only)."""
78
+ width = hidden.shape[-1]
79
+ positions = selected[:, :, None].expand(-1, -1, width)
80
+ normalized = layer.norm1(hidden)
81
+ attn = layer.self_attn
82
+ query = F.linear(normalized.gather(1, positions),
83
+ attn.in_proj_weight[:width], attn.in_proj_bias[:width])
84
+ key, value = F.linear(normalized, attn.in_proj_weight[width:],
85
+ attn.in_proj_bias[width:]).chunk(2, dim=-1)
86
+ batch = hidden.shape[0]
87
+ heads = attn.num_heads
88
+ split = lambda x: x.reshape(batch, -1, heads, width // heads).transpose(1, 2)
89
+ attended = F.scaled_dot_product_attention(
90
+ split(query), split(key), split(value),
91
+ attn_mask=attention_mask[:, None, None, :].bool(), dropout_p=0.)
92
+ attended = attended.transpose(1, 2).reshape(batch, selected.shape[1], width)
93
+ result = hidden.gather(1, positions) + attn.out_proj(attended)
94
+ return result + layer.linear2(layer.activation(layer.linear1(layer.norm2(result))))
95
+
96
+ def save_pretrained(self, directory):
97
+ root = Path(directory)
98
+ root.mkdir(parents=True, exist_ok=True)
99
+ self.encoder.config.save_pretrained(root / 'encoder')
100
+ config = dict(format_version=1, architecture='JuliaDecisionModel',
101
+ weight_dtype=str(next(self.parameters()).dtype).replace('torch.', ''), **self.settings)
102
+ (root / 'julia_config.json').write_text(json.dumps(config, indent=2) + '\n')
103
+ save_file({k: v.detach().cpu().contiguous() for k, v in self.state_dict().items()},
104
+ str(root / 'model.safetensors'), metadata={'format': 'pt', 'family': 'julia'})
105
+
106
+ @classmethod
107
+ def from_pretrained(cls, directory, *, memory_map=False):
108
+ root = Path(directory)
109
+ if (root / 'julia_config.json').exists():
110
+ config = json.loads((root / 'julia_config.json').read_text())
111
+ if config['format_version'] != 1:
112
+ raise ValueError('Unsupported Julia checkpoint format')
113
+ settings = {k: config[k] for k in ('head_layers', 'n_act', 'dropout')}
114
+ else:
115
+ config = json.loads((root / 'rl_agent_config.json').read_text())
116
+ settings = dict(head_layers=config['head_layers'], n_act=len(config.get('act_costs', {})) + 1)
117
+ encoder_config = AutoConfig.from_pretrained(root / 'encoder', trust_remote_code=False)
118
+ encoder_config.reference_compile = False
119
+ try:
120
+ from transformers.initialization import no_init_weights
121
+ except ImportError:
122
+ from transformers.modeling_utils import no_init_weights
123
+ with no_init_weights():
124
+ encoder = AutoModel.from_config(encoder_config, attn_implementation='sdpa', trust_remote_code=False)
125
+ model = cls(encoder, **settings)
126
+ if not memory_map and config.get('weight_dtype') == 'bfloat16':
127
+ model.to(dtype=torch.bfloat16)
128
+ # Inference can retain safetensors' private file-backed storage rather than
129
+ # copying the entire (mostly unused vocabulary) embedding into anonymous RAM.
130
+ # The default copy path remains available for training and mutable checkpoints.
131
+ model.load_state_dict(load_file(str(root / 'model.safetensors')),
132
+ strict=True, assign=memory_map)
133
+ return model
julia/probabilities.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Presentation rules for decision probabilities; logits and selection stay raw."""
2
+
3
+
4
+ def display_probabilities(probabilities):
5
+ values = list(probabilities)
6
+ if not values:
7
+ return values
8
+
9
+ winner = max(range(len(values)), key=values.__getitem__)
10
+ if values[winner] > 0.95 and all(
11
+ value < 0.045 for i, value in enumerate(values) if i != winner
12
+ ):
13
+ return [1.0 if i == winner else 0.0 for i in range(len(values))]
14
+
15
+ visible = [value if value >= 0.01 else 0.0 for value in values]
16
+ total = sum(visible)
17
+ return [value / total for value in visible]
julia/router/.gitignore ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ build/
2
+ __pycache__/
3
+ *.pyc
julia/router/README.md ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Julia router
2
+
3
+ Resident inference for Julia on **CPU or CUDA**, located in the existing `julia`
4
+ package. No separate repository. Python loads Bend-generated C with `ctypes`;
5
+ there is no compiler or subprocess in the request path.
6
+
7
+ ## Build and use
8
+
9
+ From the **Julia-1** directory (the outer workspace also has an older package
10
+ named `julia`, so running there can import the wrong one):
11
+
12
+ ```bash
13
+ python -m pip install -e .
14
+ python -m julia.router.build --bend /home/klee/.bend/bin/bend
15
+ ```
16
+
17
+ The native build requires Bend **2.0.27**, Clang, and Linux. It checks `native/PROOF.bend`, generated
18
+ function arities, and the borrowed-tree contract before compiling. `--native-cpu`
19
+ adds host-specific machine instructions; use that build only on compatible CPUs. The bridge
20
+ uses private Bend runtime internals, so upgrading Bend requires adapting and
21
+ retesting the bridge. `--output` selects another library location; pass that
22
+ path as `library=` or set `JULIA_ROUTER_LIBRARY` at runtime.
23
+
24
+ ```python
25
+ from julia.router import FastEngine
26
+
27
+ engine = FastEngine('/path/to/real/checkpoint', device='cpu', batch_size=16)
28
+ # device='cuda' uses CUDA/BF16 and device-side softmax/argmax.
29
+ rows = [{
30
+ 'state': 'Preciso trocar minha senha.',
31
+ 'question': 'Qual é a intenção?',
32
+ 'options': ['Redefinir senha', 'Cancelar conta', 'Consultar saldo'],
33
+ }]
34
+ print(engine.predict(rows))
35
+ print(engine.predict(rows, probabilities=False)) # return only indices
36
+ ```
37
+
38
+ Use real checkpoint files, not Git LFS pointer files. The 2026-09-23 runtime audit
39
+ loaded the actual 551 MiB checkpoint offline and exercised resident inference.
40
+
41
+ ## Faster inference path
42
+
43
+ - Bounded LRU caches reuse token fragments and full request encodings. Model
44
+ forward passes still run on **every request**; no answer cache inflates timings.
45
+ - Sort by encoded length before bounded microbatches, then restore request order.
46
+ This reduces padding on mixed-length inputs; it does not guarantee a speedup
47
+ for every batch or request size.
48
+ - Build one NumPy arena for integer inputs, with zero-copy CPU tensor views.
49
+ CUDA uses two pinned bulk transfers instead of per-field transfers.
50
+ - Softmax and argmax run on the model's device. CUDA does not copy hidden states
51
+ or attention matrices into CPU Bend kernels.
52
+ - `compile_model=True` enables optional `torch.compile(dynamic=True)`. Compilation
53
+ adds first-call cost and requires validation on the target device.
54
+ - CPU inference retains private file-backed safetensors storage instead of copying
55
+ the full vocabulary embedding into anonymous RAM. Keep checkpoint files immutable
56
+ while loaded; use `memory_map=False` for a detached copy. Training model loading
57
+ still copies by default.
58
+ - CPU PyTorch heads compute only option queries and feed-forward outputs in the
59
+ final layer. Full context keys/values are retained. `marker_only_head=False`
60
+ selects the original path; training and the separate experimental Bend normalization head use the original
61
+ path automatically. CUDA retains the original default until hardware validation.
62
+ - `strict_encoding=True` rejects marker injection and any question/option/state
63
+ truncation. `encoding_info(rows)` audits the same cached encoding used in inference;
64
+ the game worker no longer tokenizes each request twice.
65
+ - Preserve Julia's original marker serialization, truncation, weights, and
66
+ option order. Returned probabilities use the model card presentation rules;
67
+ `logits()` retains raw scores.
68
+
69
+ CPU defaults to Python/PyTorch (`torch`), with no Bend build required. CUDA also
70
+ uses PyTorch. Select `transformer_backend='bend-dense'` explicitly to use the
71
+ optional CPU FP32 Bend encoder. `compile_model=True` works with the default Torch
72
+ backend.
73
+
74
+ ## Bend transformer operations
75
+
76
+ `native/router.bend` implements candidate argmax, numerically stabilized softmax,
77
+ and two-pass LayerNorm (mean, centered variance, affine scale/bias). Balanced
78
+ `Leaf`/`Fork` trees expose candidate reductions. LayerNorm forks over independent
79
+ rows and uses flat tail loops over feature lists within each row, following the
80
+ Bend guide's coarse-work/flat-leaf cost model. All use Bend's own heap. The C adapter transports arrays and owns the ABI; it does
81
+ not duplicate the numerical algorithms.
82
+
83
+ ```python
84
+ from julia.router import BendReducer, FastEngine
85
+
86
+ bend = BendReducer()
87
+ index, probabilities = bend.softmax([1.0, 3.0, -2.0])
88
+ normalized = bend.layernorm([[1.0, 2.0, 3.0]])
89
+
90
+ # Explicit experimental CPU transformer-head backend:
91
+ engine = FastEngine('/path/to/checkpoint', device='cpu',
92
+ transformer_backend='bend', bend_postprocess=True)
93
+
94
+ # All 88 encoder projections of the real checkpoint also execute in Bend:
95
+ # Set JULIA_BEND_THREADS=4 before creating the engine.
96
+ engine = FastEngine('/path/to/checkpoint', device='cpu',
97
+ transformer_backend='bend-dense', bend_postprocess=True)
98
+ ```
99
+
100
+ The experimental head executes its pre-attention and pre-MLP LayerNorms, plus
101
+ scorer LayerNorm, in Bend. Dense projections, attention, and the encoder remain
102
+ PyTorch. This is **not a complete transformer rewrite in Bend**. Bend head math
103
+ is inference-only CPU float32. The CPU library defaults to up to eight available Bend runtime
104
+ workers. Set `JULIA_BEND_THREADS=4` **before the first native call** to test row
105
+ parallelism; the worker count is fixed for the process. Pool sizing depends on the kernel and workload; the older normalization-only
106
+ experiment does not determine the dense backend defaults. Reduction order differs from the
107
+ old tree algorithm and PyTorch: probabilities can differ by floating-point rounding,
108
+ and effectively tied choices may choose a different index.
109
+ Calls share a mutex because Bend's runtime has global state. Instantiate model
110
+ workers in spawned processes; do not fork an active inference process.
111
+
112
+ ### Resident dense projections (CPU default)
113
+
114
+ `bend-dense` executes all 88 encoder projections in Bend. Resident packed FP32
115
+ arrays replace linked weight trees. Contiguous 8×8 tiles expose vector arithmetic;
116
+ coarse parallel ranges end in flat tail loops. Shared input/weight handles are
117
+ read-only, output tiles are disjoint, and all handles are joined after evaluation.
118
+ The bridge packs/transports buffers and checks bounds; the public API validates
119
+ finite values, while engine inference validates weights once and final logits.
120
+
121
+ PyTorch retains embeddings, SDPA, elementwise activations and the selected-output
122
+ decision head. The native pool defaults to up to eight available CPUs; override
123
+ with `JULIA_BEND_THREADS` before the first call. `JULIA_BEND_TILE_GRAIN` overrides
124
+ the adaptive projection-specific chunk size. The snapshot rejects changed weights and must be recreated after mutation.
125
+ Use the training loader to save or train checkpoints, not an installed backend.
126
+
127
+ Packed kernels use
128
+ validated unsafe array sharing and bounded loops; the tree shape proofs do not
129
+ constitute a formal proof of packed-array memory safety or floating-point math.
130
+
131
+ ## Larger choice sets
132
+
133
+ ```python
134
+ from julia.router import Router, FastEngine
135
+ router = Router(FastEngine('/path/to/checkpoint', device='cuda'), survivors=2)
136
+ result = router.route(row_with_up_to_4096_options)
137
+ ```
138
+
139
+ Julia's trained head still supports **2–20** options. Larger *choice* requests
140
+ use batched groups and rerank survivors until a final group remains. A group
141
+ retains only its winner when raw softmax gives it over 95% and every other
142
+ option is below 4.5%; otherwise it retains the configured survivor count.
143
+ This can reduce later model calls for decisive groups, but group probabilities
144
+ are not comparable across different groups. This adds model calls and can discard the correct candidate;
145
+ it is a capacity feature, not a speed or quality guarantee. Final probabilities
146
+ are conditional on `result.candidates`, never a fabricated global distribution.
147
+ `model_rows` and `cache_hits` report aggregate work for the whole `route_many`
148
+ call. Optional `cache_size` caches logits; it defaults to zero. Clear caches
149
+ with `clear_cache()` after changing weights, tokenization or inference settings.
150
+
151
+ ## Tests
152
+
153
+ ```bash
154
+ python -m unittest discover -s julia/router/tests -v
155
+ ```
156
+
157
+ The tests exercise CPU inference and optional native behavior where available.
julia/router/__init__.py ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Python-compatible Julia router powered by compiled Bend trees."""
2
+ from .native import BendReducer
3
+ from .router import Router, RouteResult
4
+
5
+ __all__ = ['BendReducer', 'Router', 'RouteResult', 'FastEngine']
6
+
7
+
8
+ def __getattr__(name):
9
+ if name == 'FastEngine':
10
+ from .engine import FastEngine
11
+ return FastEngine
12
+ raise AttributeError(name)
julia/router/build.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Explicit, reproducible native build: python -m julia.router.build."""
2
+ import argparse
3
+ import os
4
+ import re
5
+ from pathlib import Path
6
+ import shutil
7
+ import subprocess
8
+ import tempfile
9
+
10
+
11
+ def build(output=None, bend=None, native_cpu=False):
12
+ source = Path(__file__).parent / 'native' / 'router.bend'
13
+ bend = bend or shutil.which('bend') or str(Path.home() / '.bend/bin/bend')
14
+ version = subprocess.check_output([bend, 'version'], text=True).strip()
15
+ if not re.search(r'(?<![0-9.])2\.0\.27(?![0-9.])', version):
16
+ raise RuntimeError(f'Native adapter requires Bend 2.0.27; found {version}')
17
+ subprocess.run([bend, str(source.with_name('PROOF.bend'))], check=True)
18
+ output = Path(output or Path(__file__).parent / 'build/libjulia_router.so').resolve()
19
+ output.parent.mkdir(parents=True, exist_ok=True)
20
+ with tempfile.TemporaryDirectory(dir=output.parent) as temp:
21
+ generated = Path(temp) / 'router.c'
22
+ library = Path(temp) / 'router.so'
23
+ subprocess.run([bend, str(source), '-o', str(generated)], check=True)
24
+ emitted = generated.read_text()
25
+ if not all(re.search(pattern, emitted) for pattern in (
26
+ r'#define CUBE_T\s+128', r'#define LINE\s+16',
27
+ r'static u32\s+CUBE_LOG = 7;', r'#define RING_LOG\s+\(17 - CUBE_LOG\)')):
28
+ raise RuntimeError('CPU scheduling layout changed; refusing unsafe native build')
29
+ # The adapter depends on these private generated ABI/ownership contracts.
30
+ for name, arity in [('WINNER', 1), ('SOFTMAX', 1), ('LAYERNORM', 3), ('NORM_ROWS', 3), ('DENSE_TILE', 2), ('DENSE_BATCH', 2), ('PACKED_GEMM', 9)]:
31
+ match = re.search(rf'#define FID_{name} (\d+)', emitted)
32
+ table = re.search(r'FID_ARITY_T\[\] = \{([^}]+)', emitted)
33
+ if not match or not table or int(table[1].split(',')[int(match[1])]) != arity:
34
+ raise RuntimeError(f'Unsupported generated Bend ABI for {name}')
35
+ winner = emitted.split('WL_CASE(FID_WINNER)', 1)[1].split('#endif', 1)[0]
36
+ if 'term_peek(e, _tree_0)' not in winner:
37
+ raise RuntimeError('WINNER ownership changed; refusing unsafe native build')
38
+ tile = emitted.split('WL_CASE(FID_DENSE_TILE)', 1)[1].split('#endif', 1)[0]
39
+ if 'term_peek(e, _weights_0)' not in tile or 'term_peek(e, _ws_0)' not in emitted:
40
+ raise RuntimeError('Tiled weights must be borrowed; refusing unsafe native build')
41
+ batch = emitted.split('WL_CASE(FID_DENSE_BATCH)', 1)[1].split('#endif', 1)[0]
42
+ if 'term_peek(e, _tiles_0)' not in batch:
43
+ raise RuntimeError('Dense input tiles must be borrowed; refusing unsafe native build')
44
+ packed = emitted.split('WL_CASE(FID_PACKED_DOT)', 1)[1].split('#endif', 1)[0]
45
+ if not all(fragment in packed for fragment in ('r0 = _x_0;', 'r1 = _w_0;', 'r2 = _o_0;', 'WL_RETN(3);')):
46
+ raise RuntimeError('Packed result ABI changed; refusing unsafe native build')
47
+ from .specialize import packed_cpu
48
+ generated.write_text(packed_cpu(emitted))
49
+ subprocess.run([os.environ.get('CC', 'clang'), '-O3', '-ffp-contract=fast',
50
+ *(['-march=native'] if native_cpu else []), '-fPIC', '-shared',
51
+ '-pthread', str(generated), '-lm', '-ldl', '-o', str(library)], check=True)
52
+ os.replace(library, output)
53
+ return output
54
+
55
+
56
+ if __name__ == '__main__':
57
+ parser = argparse.ArgumentParser()
58
+ parser.add_argument('--output')
59
+ parser.add_argument('--bend')
60
+ parser.add_argument('--native-cpu', action='store_true', help='Optimize for this CPU; do not distribute to other CPUs')
61
+ args = parser.parse_args()
62
+ print(build(args.output, args.bend, args.native_cpu))
julia/router/encoder.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Inference-only ModernBERT path for decision models (no unused outputs)."""
2
+ from types import MethodType
3
+ import torch
4
+ from transformers.modeling_outputs import BaseModelOutput
5
+
6
+
7
+ def _decision_forward(self, input_ids=None, attention_mask=None, **kwargs):
8
+ if self.training or kwargs or input_ids is None or attention_mask is None:
9
+ return self._julia_original_forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
10
+ position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0)
11
+ full_mask, local_mask = self._update_attention_mask(attention_mask, output_attentions=False)
12
+ hidden = self.embeddings(input_ids=input_ids)
13
+ # Upstream iterates config.layer_types (one entry per layer), overwriting
14
+ # the same two dictionary entries 22 times in this checkpoint.
15
+ positions = {kind: self.rotary_emb(hidden, position_ids, kind) for kind in self._julia_attention_types}
16
+ for layer in self.layers:
17
+ hidden = layer(hidden, attention_mask=full_mask, sliding_window_mask=local_mask,
18
+ position_ids=position_ids, cu_seqlens=None, max_seqlen=None,
19
+ position_embeddings=positions[layer.attention_type], output_attentions=False)[0]
20
+ return BaseModelOutput(last_hidden_state=self.final_norm(hidden))
21
+
22
+
23
+ def specialize_decision_encoder(model):
24
+ encoder = model.encoder
25
+ if encoder.config.model_type != 'modernbert' or encoder.config._attn_implementation != 'sdpa':
26
+ return False
27
+ if hasattr(encoder, '_julia_original_forward'):
28
+ return True
29
+ encoder._julia_original_forward = encoder.forward
30
+ encoder._julia_attention_types = tuple(dict.fromkeys(encoder.config.layer_types))
31
+ encoder.forward = MethodType(_decision_forward, encoder)
32
+ return True
julia/router/engine.py ADDED
@@ -0,0 +1,225 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Resident CPU/CUDA inference with cached encoding and length-aware batches."""
2
+ from collections import OrderedDict
3
+ import json
4
+ from pathlib import Path
5
+ import threading
6
+
7
+ import numpy as np
8
+ import torch
9
+
10
+ from ..data import sequence, validate_row
11
+ from ..probabilities import display_probabilities
12
+ from ..inference import TransformerEngine as Engine
13
+
14
+
15
+ class _TokenCache:
16
+ def __init__(self, tokenizer, capacity):
17
+ self.tokenizer, self.capacity = tokenizer, capacity
18
+ self.cache = OrderedDict()
19
+
20
+ def __getattr__(self, name):
21
+ return getattr(self.tokenizer, name)
22
+
23
+ def __call__(self, text, add_special_tokens=False):
24
+ if text not in self.cache:
25
+ ids = self.tokenizer(text, add_special_tokens=False)['input_ids']
26
+ if self.capacity:
27
+ self.cache[text] = tuple(ids)
28
+ while len(self.cache) > self.capacity:
29
+ self.cache.popitem(last=False)
30
+ return {'input_ids': ids}
31
+ self.cache.move_to_end(text)
32
+ return {'input_ids': list(self.cache[text])}
33
+
34
+
35
+ class FastEngine(Engine):
36
+ """Same weights/serialization as Engine; no cached model predictions.
37
+
38
+ Optional torch.compile specializes the transformer, while Bend handles CPU
39
+ softmax/selection via ctypes. CUDA softmax stays on-device to avoid a roundtrip.
40
+ """
41
+ def __init__(self, checkpoint, device='cpu', max_length=1024, head_length=256,
42
+ batch_size=16, encoding_cache=2048, token_cache=8192,
43
+ compile_model=False, library=None, bend_postprocess=False, transformer_backend=None,
44
+ strict_encoding=False, marker_only_head=None, memory_map=True, padding_ratio=1.25):
45
+ if not 1 <= padding_ratio <= 16:
46
+ raise ValueError('padding_ratio must be between 1 and 16')
47
+ self.padding_ratio = padding_ratio
48
+ if batch_size < 1 or encoding_cache < 0 or token_cache < 0:
49
+ raise ValueError('Invalid batch/cache size')
50
+ if transformer_backend not in (None, 'torch', 'bend', 'bend-dense'):
51
+ raise ValueError('transformer_backend must be torch, bend, or bend-dense')
52
+ if transformer_backend in ('bend', 'bend-dense') and (str(device) != 'cpu' or compile_model):
53
+ raise ValueError('Bend requires eager CPU inference; compilation requires the torch backend')
54
+ weights = Path(checkpoint) / 'model.safetensors'
55
+ if weights.is_file():
56
+ with weights.open('rb') as stream:
57
+ if stream.read(80).startswith(b'version https://git-lfs.github.com/spec/v1'):
58
+ raise ValueError('Checkpoint contains Git LFS pointers; fetch the real model weights first')
59
+ super().__init__(checkpoint, device, max_length, head_length, memory_map=memory_map)
60
+ self.model.marker_only_head = (self.device.type == 'cpu' if marker_only_head is None else marker_only_head)
61
+ self.strict_encoding = strict_encoding
62
+ self.batch_size = batch_size
63
+ self.max_length, self.head_length = max_length, head_length
64
+ self.encoding_cache = encoding_cache
65
+ self._encoded = OrderedDict()
66
+ self._tokens = _TokenCache(self.tokenizer, token_cache)
67
+ self._lock = threading.RLock()
68
+ self.bend = None
69
+ if bend_postprocess and self.device.type == 'cpu':
70
+ from .native import BendReducer
71
+ self.bend = BendReducer(library)
72
+ if transformer_backend is None:
73
+ transformer_backend = 'torch'
74
+ self.transformer_backend = transformer_backend
75
+ if transformer_backend in ('bend', 'bend-dense'):
76
+ if next(self.model.parameters()).dtype != torch.float32:
77
+ raise ValueError('Bend backends require FP32 weights')
78
+ from .native import BendReducer
79
+ from .transformer import install_bend_head
80
+ reducer = self.bend or BendReducer(library)
81
+ if transformer_backend == 'bend':
82
+ install_bend_head(self.model, reducer)
83
+ if transformer_backend == 'bend-dense':
84
+ from .transformer import install_bend_encoder
85
+ self.bend_projection_count = install_bend_encoder(self.model, reducer)
86
+ from .encoder import specialize_decision_encoder
87
+ self.encoder_specialized = specialize_decision_encoder(self.model)
88
+ self.forward = self.model
89
+ if compile_model:
90
+ self.forward = torch.compile(self.model, dynamic=True)
91
+
92
+ def clear_cache(self):
93
+ with self._lock:
94
+ self._encoded.clear()
95
+ self._tokens.cache.clear()
96
+
97
+ def _encode(self, rows):
98
+ result = []
99
+ for i, row in enumerate(rows):
100
+ validate_row(row, i + 1)
101
+ key = json.dumps([self.max_length, self.head_length, self.strict_encoding,
102
+ row['state'], row['question'], row['options'],
103
+ row.get('type', 'choice')], ensure_ascii=False, allow_nan=False)
104
+ encoded = self._encoded.get(key)
105
+ if encoded is None:
106
+ encoded = sequence(self._tokens, row, self.max_length, self.head_length,
107
+ strict=self.strict_encoding)
108
+ if self.encoding_cache:
109
+ self._encoded[key] = encoded
110
+ while len(self._encoded) > self.encoding_cache:
111
+ self._encoded.popitem(last=False)
112
+ else:
113
+ self._encoded.move_to_end(key)
114
+ result.append(encoded)
115
+ return result
116
+
117
+ def encoding_info(self, rows):
118
+ """Audit the same cached encoding used by inference, without retokenizing."""
119
+ if not self.strict_encoding:
120
+ raise ValueError('Lossless encoding audit requires strict_encoding=True')
121
+ with self._lock:
122
+ return [dict(tokens=len(item['ids']), optionTokens=list(item['option_tokens']),
123
+ headLength=self.head_length, stateTruncated=False, optionsTruncated=False)
124
+ for item in self._encode(rows)]
125
+
126
+ def _pack(self, encoded):
127
+ length = min(self.max_length, (max(len(x['ids']) for x in encoded) + 7) // 8 * 8)
128
+ count = max(len(x['markers']) for x in encoded)
129
+ size = len(encoded)
130
+ arena = np.zeros(size * (2 * length + count + 1), dtype=np.int64)
131
+ end = size * length
132
+ ids = arena[:end].reshape(size, length)
133
+ ids.fill(self.tokenizer.pad_token_id)
134
+ attention = arena[end:2 * end].reshape(size, length)
135
+ positions = arena[2 * end:2 * end + size * count].reshape(size, count)
136
+ qtype = arena[-size:]
137
+ mask = np.zeros((size, count), dtype=np.bool_)
138
+ for i, item in enumerate(encoded):
139
+ n, k = len(item['ids']), len(item['markers'])
140
+ ids[i, :n] = item['ids']
141
+ attention[i, :n] = 1
142
+ positions[i, :k] = item['markers']
143
+ mask[i, :k] = True
144
+ qtype[i] = item['qtype']
145
+ host = torch.from_numpy(arena)
146
+ marker_mask = torch.from_numpy(mask)
147
+ if self.device.type == 'cuda':
148
+ host = host.pin_memory().to(self.device, non_blocking=True)
149
+ marker_mask = marker_mask.pin_memory().to(self.device, non_blocking=True)
150
+ return dict(input_ids=host[:end].view(size, length),
151
+ attention_mask=host[end:2 * end].view(size, length),
152
+ marker_pos=host[2 * end:2 * end + size * count].view(size, count),
153
+ qtype=host[-size:], marker_mask=marker_mask)
154
+
155
+ def _batch_indices(self, encoded):
156
+ # Bound padding inflation, not just batch count: a single long request
157
+ # must not make every short request run the entire encoder at its length.
158
+ order = sorted(range(len(encoded)), key=lambda i: len(encoded[i]['ids']))
159
+ group, tokens = [], 0
160
+ for index in order:
161
+ length = len(encoded[index]['ids'])
162
+ if group and (len(group) == self.batch_size or
163
+ (len(group) + 1) * length > self.padding_ratio * (tokens + length)):
164
+ yield group
165
+ group, tokens = [], 0
166
+ group.append(index)
167
+ tokens += length
168
+ if group:
169
+ yield group
170
+
171
+ def _batches(self, encoded):
172
+ for indices in self._batch_indices(encoded):
173
+ batch = self._pack([encoded[i] for i in indices])
174
+ with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16,
175
+ enabled=self.device.type == 'cuda'):
176
+ values = self.forward(**batch)
177
+ # Padding must never win against real very negative logits.
178
+ values = values.masked_fill(~batch['marker_mask'], -torch.inf)
179
+ if not (torch.isfinite(values) | ~batch['marker_mask']).all():
180
+ raise FloatingPointError('Inference returned nonfinite logits')
181
+ yield indices, values
182
+
183
+ @torch.inference_mode()
184
+ def logits(self, rows):
185
+ with self._lock:
186
+ encoded = self._encode(rows)
187
+ result = [None] * len(rows)
188
+ for indices, values in self._batches(encoded):
189
+ host = values.cpu().tolist()
190
+ for i, scores in zip(indices, host):
191
+ result[i] = scores[:len(encoded[i]['markers'])]
192
+ return result
193
+
194
+ @torch.inference_mode()
195
+ def predict(self, rows, *, probabilities=True):
196
+ with self._lock:
197
+ encoded = self._encode(rows)
198
+ result = [None] * len(rows)
199
+ for indices, values in self._batches(encoded):
200
+ if self.bend is not None:
201
+ host = values.contiguous().numpy()
202
+ for local, i in enumerate(indices):
203
+ scores = host[local, :len(encoded[i]['markers'])]
204
+ if probabilities:
205
+ best, probs = self.bend.softmax(scores)
206
+ else:
207
+ best = self.bend.argmax(scores)
208
+ result[i] = dict(index=best)
209
+ if probabilities:
210
+ result[i]['probabilities'] = display_probabilities(probs.tolist())
211
+ else:
212
+ best_device = values.argmax(-1)
213
+ if probabilities and self.device.type == 'cuda':
214
+ host = torch.cat((best_device[:, None].to(values.dtype),
215
+ values.softmax(-1)), dim=1).cpu().tolist()
216
+ best = [int(row[0]) for row in host]
217
+ probs = [row[1:] for row in host]
218
+ else:
219
+ best = best_device.cpu().tolist()
220
+ probs = values.softmax(-1).tolist() if probabilities else None
221
+ for local, i in enumerate(indices):
222
+ result[i] = dict(index=best[local])
223
+ if probabilities:
224
+ result[i]['probabilities'] = display_probabilities(probs[local][:len(encoded[i]['markers'])])
225
+ return result
julia/router/native.py ADDED
@@ -0,0 +1,159 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ctypes boundary; no subprocess or compilation in the request path."""
2
+ import ctypes
3
+ import math
4
+ import os
5
+ from pathlib import Path
6
+
7
+
8
+ class BendReducer:
9
+ def __init__(self, library=None):
10
+ path = Path(library or os.environ.get('JULIA_ROUTER_LIBRARY') or
11
+ Path(__file__).parent / 'build/libjulia_router.so')
12
+ try:
13
+ self._lib = ctypes.CDLL(str(path.resolve()))
14
+ except OSError as error:
15
+ raise RuntimeError(
16
+ f'Cannot load Bend runtime {path.resolve()}. From Julia-1 run '
17
+ 'python -m julia.router.build --native-cpu; or explicitly select '
18
+ 'transformer_backend="torch".') from error
19
+ self._argmax = self._lib.julia_router_argmax
20
+ self._argmax.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.c_uint32,
21
+ ctypes.POINTER(ctypes.c_uint32)]
22
+ self._argmax.restype = ctypes.c_int
23
+ self._ptr = ctypes.POINTER(ctypes.c_float)
24
+ self._softmax = self._lib.julia_router_softmax
25
+ self._softmax.argtypes = [self._ptr, ctypes.c_uint32, self._ptr, ctypes.POINTER(ctypes.c_uint32)]
26
+ self._softmax.restype = ctypes.c_int
27
+ self._layernorm = self._lib.julia_router_layernorm
28
+ self._layernorm.argtypes = [self._ptr, self._ptr, self._ptr, ctypes.c_uint32,
29
+ ctypes.c_uint32, ctypes.c_float, self._ptr]
30
+ self._layernorm.restype = ctypes.c_int
31
+
32
+ def profile(self, *, enabled=False, reset=False):
33
+ """Process-wide packed projection timings; opt in only for diagnostics."""
34
+ fn = self._lib.julia_router_profile
35
+ fn.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double)]
36
+ fn.restype = None
37
+ values = (ctypes.c_double * 4)()
38
+ fn(enabled, reset, values)
39
+ return dict(zip(('pack_ms', 'execute_ms', 'unpack_ms', 'calls'), values))
40
+
41
+ def argmax(self, scores):
42
+ scores = tuple(scores)
43
+ if not 1 <= len(scores) <= 1048576:
44
+ raise ValueError('Expected 1–1048576 scores')
45
+ if not all(math.isfinite(x) for x in scores):
46
+ raise ValueError('Scores must be finite')
47
+ values = (ctypes.c_float * len(scores))(*scores)
48
+ result = ctypes.c_uint32()
49
+ code = self._argmax(values, len(scores), ctypes.byref(result))
50
+ if code:
51
+ raise ValueError(f'Bend rejected scores (status {code})')
52
+ return result.value
53
+
54
+ def softmax(self, scores):
55
+ import numpy as np
56
+ values = np.ascontiguousarray(scores, dtype=np.float32)
57
+ if values.ndim != 1 or not 1 <= values.size <= 1048576:
58
+ raise ValueError('Expected a nonempty vector of at most 1048576 scores')
59
+ output = np.empty_like(values)
60
+ index = ctypes.c_uint32()
61
+ code = self._softmax(values.ctypes.data_as(self._ptr), values.size,
62
+ output.ctypes.data_as(self._ptr), ctypes.byref(index))
63
+ if code:
64
+ raise ValueError(f'Bend rejected scores (status {code})')
65
+ return index.value, output
66
+
67
+ def layernorm(self, values, gamma=None, beta=None, epsilon=1e-5):
68
+ import numpy as np
69
+ x = np.ascontiguousarray(values, dtype=np.float32)
70
+ if x.ndim < 1 or x.size == 0:
71
+ raise ValueError('Expected a nonempty feature array')
72
+ width = x.shape[-1]
73
+ g = np.ones(width, np.float32) if gamma is None else np.ascontiguousarray(gamma, dtype=np.float32)
74
+ b = np.zeros(width, np.float32) if beta is None else np.ascontiguousarray(beta, dtype=np.float32)
75
+ if g.shape != (width,) or b.shape != (width,):
76
+ raise ValueError('gamma and beta must match the last dimension')
77
+ if x.size > 1048576 or width > 65536:
78
+ raise ValueError('LayerNorm input exceeds native capacity')
79
+ output = np.empty_like(x)
80
+ code = self._layernorm(x.ctypes.data_as(self._ptr), g.ctypes.data_as(self._ptr),
81
+ b.ctypes.data_as(self._ptr), x.size // width, width, epsilon,
82
+ output.ctypes.data_as(self._ptr))
83
+ if code:
84
+ raise ValueError(f'Bend rejected LayerNorm inputs (status {code})')
85
+ return output
86
+
87
+
88
+ class BendMatrix:
89
+ """Resident immutable FP32 projection weights owned by the Bend heap."""
90
+ def __init__(self, weights, reducer=None):
91
+ import numpy as np
92
+ import threading
93
+ import weakref
94
+ self.reducer = reducer or BendReducer()
95
+ self._lock = threading.RLock()
96
+ values = np.ascontiguousarray(weights, dtype=np.float32)
97
+ if values.ndim != 2 or not all(values.shape):
98
+ raise ValueError('Projection weights must be a nonempty matrix')
99
+ self.outputs, self.inputs = values.shape
100
+ lib = self.reducer._lib
101
+ create = lib.julia_router_matrix_create
102
+ create.argtypes = [self.reducer._ptr, ctypes.c_uint32, ctypes.c_uint32]
103
+ create.restype = ctypes.c_void_p
104
+ release = lib.julia_router_matrix_free
105
+ release.argtypes = [ctypes.c_void_p]
106
+ release.restype = None
107
+ self._linear = lib.julia_router_linear
108
+ self._linear.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p]
109
+ self._linear.restype = ctypes.c_int
110
+ self._linear_inference = lib.julia_router_linear_inference
111
+ self._linear_inference.argtypes = self._linear.argtypes
112
+ self._linear_inference.restype = ctypes.c_int
113
+ self._handle = create(values.ctypes.data_as(self.reducer._ptr), self.inputs, self.outputs)
114
+ if not self._handle:
115
+ raise ValueError('Bend rejected projection dimensions or nonfinite weights')
116
+ self._finalize = weakref.finalize(self, release, self._handle)
117
+
118
+ def close(self):
119
+ with self._lock:
120
+ self._finalize()
121
+
122
+ def tensor(self, values, *, validate=True):
123
+ """Fill a contiguous Torch output directly, without NumPy/list/concat wrappers."""
124
+ import torch
125
+ if values.device.type != 'cpu' or values.dtype != torch.float32 or values.shape[-1] != self.inputs:
126
+ raise ValueError('Bend tensor projection requires matching CPU FP32 inputs')
127
+ x = values.contiguous()
128
+ rows = x.numel() // self.inputs
129
+ if not rows:
130
+ raise ValueError('Projection input must be nonempty')
131
+ out = torch.empty((*x.shape[:-1], self.outputs), dtype=x.dtype, device=x.device)
132
+ linear = self._linear if validate else self._linear_inference
133
+ limit = min(1048576 // self.inputs, 1048576 // self.outputs)
134
+ with self._lock:
135
+ if not self._finalize.alive:
136
+ raise RuntimeError('Bend matrix has been closed')
137
+ for start in range(0, rows, limit):
138
+ code = linear(self._handle,
139
+ x.data_ptr() + start * self.inputs * 4,
140
+ min(limit, rows - start),
141
+ out.data_ptr() + start * self.outputs * 4)
142
+ if code:
143
+ raise ValueError(f'Bend projection failed (status {code})')
144
+ return out
145
+
146
+ def __call__(self, values):
147
+ import numpy as np
148
+ x = np.ascontiguousarray(values, dtype=np.float32)
149
+ if x.ndim < 1 or x.shape[-1] != self.inputs or not x.size:
150
+ raise ValueError('Projection input must match resident weights')
151
+ out = np.empty((*x.shape[:-1], self.outputs), dtype=np.float32)
152
+ with self._lock:
153
+ if not self._finalize.alive:
154
+ raise RuntimeError('Bend matrix has been closed')
155
+ code = self._linear(self._handle, x.ctypes.data_as(self.reducer._ptr),
156
+ x.size // self.inputs, out.ctypes.data_as(self.reducer._ptr))
157
+ if code:
158
+ raise ValueError(f'Bend projection failed (status {code})')
159
+ return out
julia/router/native/LAWS.bend ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import Base
2
+ import ./router.bend as R
3
+
4
+ # Runtime ABI invariant: each output leaf has exactly one 4x7 tile.
5
+ def output_tiles(result: R.TileResult) -> Nat:
6
+ match result:
7
+ case R.JTLeaf{_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _}:
8
+ 1n
9
+ case R.JTFork{left, right}:
10
+ Nat.add(output_tiles(left), output_tiles(right))
11
+
12
+ def matrix_tiles(weights: R.JuliaMatrix) -> Nat:
13
+ match weights:
14
+ case R.JMLeaf{ws}:
15
+ 1n
16
+ case R.JMFork{left, right}:
17
+ Nat.add(matrix_tiles(left), matrix_tiles(right))
18
+
19
+ # No flat dot-product branch may omit or duplicate an output tile.
20
+ law dot_tile_shape:
21
+ for +xs: R.TileInput
22
+ for ws: R.DenseWeights
23
+ for s0: F32
24
+ for s1: F32
25
+ for s2: F32
26
+ for s3: F32
27
+ for s4: F32
28
+ for s5: F32
29
+ for s6: F32
30
+ for s7: F32
31
+ for s8: F32
32
+ for s9: F32
33
+ for s10: F32
34
+ for s11: F32
35
+ for s12: F32
36
+ for s13: F32
37
+ for s14: F32
38
+ for s15: F32
39
+ for s16: F32
40
+ for s17: F32
41
+ for s18: F32
42
+ for s19: F32
43
+ for s20: F32
44
+ for s21: F32
45
+ for s22: F32
46
+ for s23: F32
47
+ for s24: F32
48
+ for s25: F32
49
+ for s26: F32
50
+ for s27: F32
51
+ {output_tiles(R.dot_tile(xs, ws, s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27)) == 1n : Nat}
52
+
53
+ # Fork/join must preserve the output tile count for every matrix/input tree.
54
+ law dense_shape:
55
+ for +xs: R.TileInput
56
+ for weights: R.JuliaMatrix
57
+ {output_tiles(R.dense_tile(xs, weights)) == matrix_tiles(weights) : Nat}
julia/router/native/PROOF.bend ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import Base
2
+ import ./router.bend as R
3
+ import ./LAWS.bend as Laws
4
+
5
+ def Laws.dot_tile_shape(xs, ws, s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27):
6
+ match xs ws:
7
+ case R.JTVal{+x0, +x1, +x2, +x3, xt} R.JWVal{+w0, +w1, +w2, +w3, +w4, +w5, +w6, wt}:
8
+ Laws.dot_tile_shape(xt, wt, F32.add(s0, F32.mul(x0, w0)),
9
+ F32.add(s1, F32.mul(x0, w1)),
10
+ F32.add(s2, F32.mul(x0, w2)),
11
+ F32.add(s3, F32.mul(x0, w3)),
12
+ F32.add(s4, F32.mul(x0, w4)),
13
+ F32.add(s5, F32.mul(x0, w5)),
14
+ F32.add(s6, F32.mul(x0, w6)),
15
+ F32.add(s7, F32.mul(x1, w0)),
16
+ F32.add(s8, F32.mul(x1, w1)),
17
+ F32.add(s9, F32.mul(x1, w2)),
18
+ F32.add(s10, F32.mul(x1, w3)),
19
+ F32.add(s11, F32.mul(x1, w4)),
20
+ F32.add(s12, F32.mul(x1, w5)),
21
+ F32.add(s13, F32.mul(x1, w6)),
22
+ F32.add(s14, F32.mul(x2, w0)),
23
+ F32.add(s15, F32.mul(x2, w1)),
24
+ F32.add(s16, F32.mul(x2, w2)),
25
+ F32.add(s17, F32.mul(x2, w3)),
26
+ F32.add(s18, F32.mul(x2, w4)),
27
+ F32.add(s19, F32.mul(x2, w5)),
28
+ F32.add(s20, F32.mul(x2, w6)),
29
+ F32.add(s21, F32.mul(x3, w0)),
30
+ F32.add(s22, F32.mul(x3, w1)),
31
+ F32.add(s23, F32.mul(x3, w2)),
32
+ F32.add(s24, F32.mul(x3, w3)),
33
+ F32.add(s25, F32.mul(x3, w4)),
34
+ F32.add(s26, F32.mul(x3, w5)),
35
+ F32.add(s27, F32.mul(x3, w6)))
36
+ case R.JTVal{_, _, _, _, _} R.JWEnd{}:
37
+ {==}
38
+ case R.JTEnd{} R.JWVal{_, _, _, _, _, _, _, _}:
39
+ {==}
40
+ case R.JTEnd{} R.JWEnd{}:
41
+ {==}
42
+
43
+ def Laws.dense_shape(xs, weights):
44
+ match weights:
45
+ case R.JMLeaf{ws}:
46
+ Laws.dot_tile_shape(xs, ws, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
47
+ case R.JMFork{+left, +right}:
48
+ %Laws.dense_shape(xs, left) : {Nat.add(Laws.output_tiles(R.dense_tile(xs, left)), Laws.output_tiles(R.dense_tile(xs, right))) == Nat.add(_, Laws.matrix_tiles(right)) : Nat}
49
+ %Laws.dense_shape(xs, right) : {Nat.add(Laws.output_tiles(R.dense_tile(xs, left)), Laws.output_tiles(R.dense_tile(xs, right))) == Nat.add(Laws.output_tiles(R.dense_tile(xs, left)), _) : Nat}
50
+ {==}
julia/router/native/bridge.c ADDED
@@ -0,0 +1,452 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* Bend 2.0.27 ABI adapter. Arithmetic lives in router.bend.
2
+ * WINNER borrows its tree; SOFTMAX/LAYERNORM consume it (checked at build).
3
+ * The generated runtime uses process globals, hence the single shared mutex.
4
+ */
5
+ #include <pthread.h>
6
+ #if defined(__AVX2__)
7
+ #include <immintrin.h>
8
+ #endif
9
+ static pthread_mutex_t router_mutex = PTHREAD_MUTEX_INITIALIZER;
10
+ static Corpus router_heap;
11
+ /* Opt-in diagnostics, disabled during ordinary inference. Protected by router_mutex. */
12
+ static bool router_profile_enabled = false;
13
+ static double router_profile_totals[4] = {0}; /* pack, eval, unpack milliseconds; calls */
14
+ static double router_clock_ms(void) {
15
+ struct timespec now;
16
+ clock_gettime(CLOCK_MONOTONIC, &now);
17
+ return (double)now.tv_sec * 1000.0 + (double)now.tv_nsec / 1000000.0;
18
+ }
19
+ void julia_router_profile(int enabled, int reset, double *out) {
20
+ pthread_mutex_lock(&router_mutex);
21
+ if (out) memcpy(out, router_profile_totals, sizeof(router_profile_totals));
22
+ if (reset) memset(router_profile_totals, 0, sizeof(router_profile_totals));
23
+ router_profile_enabled = enabled != 0;
24
+ pthread_mutex_unlock(&router_mutex);
25
+ }
26
+ static Env router_env(void) {
27
+ if (!router_heap) {
28
+ const char *setting = getenv("JULIA_BEND_THREADS");
29
+ long available = cpu_count();
30
+ long threads = setting ? strtol(setting, NULL, 10) : (available < 8 ? available : 8);
31
+ if (threads < 1 || threads > 32) threads = 1;
32
+ /* CPU projections need about two coarse tasks per worker, not the GPU-sized
33
+ * default 16,384-lane frontier. A smaller bag keeps the same total ring
34
+ * capacity (deeper queues) and cuts empty-row scans on every pool turn. */
35
+ CUBE_LOG = 0;
36
+ while ((CUBE_T / LINE) * (1u << CUBE_LOG) < (u32)threads) ++CUBE_LOG;
37
+ const char *cube = getenv("JULIA_BEND_CUBE_LOG");
38
+ if (cube) {
39
+ char *end;
40
+ long value = strtol(cube, &end, 10);
41
+ if (*cube && !*end && value >= (long)CUBE_LOG && value <= 7) CUBE_LOG = value;
42
+ }
43
+ router_heap = corpus_setup(false, threads, 0);
44
+ io_stk = pool_stack();
45
+ }
46
+ return (Env){router_heap, ALC[0]};
47
+ }
48
+ static Term router_tree(Env e, const float *scores, u32 start, u32 end) {
49
+ if (end - start == 1)
50
+ return io_node(e, CID_LEAF, start, f32_rewrap(scores[start]));
51
+ u32 mid = start + (end - start) / 2;
52
+ Term left = router_tree(e, scores, start, mid);
53
+ Term right = router_tree(e, scores, mid, end);
54
+ return io_node(e, CID_FORK, left, right);
55
+ }
56
+ static Term router_eval(Env e, u32 fid, Term tree) {
57
+ Loc task = task_node(e, fid, TERM_HOLE, 0, 0);
58
+ e.mem[task] = tree;
59
+ return corpus_eval(e.mem, term_tsk(fid, task));
60
+ }
61
+ static void router_unpack(Env e, Term tree, float *output) {
62
+ Loc node = term_loc(tree);
63
+ if (term_aux(tree) == CID_LEAF) {
64
+ if (output) output[(u32)e.mem[node]] = f32_unbox(e.mem[node + 1]);
65
+ } else {
66
+ router_unpack(e, e.mem[node], output);
67
+ router_unpack(e, e.mem[node + 1], output);
68
+ }
69
+ heap_free(e, cls_fit(2), node);
70
+ }
71
+ static int router_valid(const float *scores, u32 count) {
72
+ if (!scores || !count || count > 1048576) return -1;
73
+ for (u32 i = 0; i < count; ++i) if (!isfinite(scores[i])) return -2;
74
+ return 0;
75
+ }
76
+ int julia_router_argmax(const float *scores, uint32_t count, uint32_t *index) {
77
+ int status = router_valid(scores, count);
78
+ if (status || !index) return status ? status : -1;
79
+ pthread_mutex_lock(&router_mutex);
80
+ Env e = router_env();
81
+ Term tree = router_tree(e, scores, 0, count);
82
+ Term result = router_eval(e, FID_WINNER, tree);
83
+ *index = (u32)e.mem[term_loc(result)];
84
+ term_drop(e, result);
85
+ router_unpack(e, tree, NULL);
86
+ pthread_mutex_unlock(&router_mutex);
87
+ return 0;
88
+ }
89
+ int julia_router_softmax(const float *scores, uint32_t count, float *out, uint32_t *index) {
90
+ int status = router_valid(scores, count);
91
+ if (status || !out || !index) return status ? status : -1;
92
+ pthread_mutex_lock(&router_mutex);
93
+ Env e = router_env();
94
+ Term tree = router_tree(e, scores, 0, count);
95
+ Term best = router_eval(e, FID_WINNER, tree);
96
+ *index = (u32)e.mem[term_loc(best)];
97
+ term_drop(e, best);
98
+ router_unpack(e, router_eval(e, FID_SOFTMAX, tree), out);
99
+ pthread_mutex_unlock(&router_mutex);
100
+ return router_valid(out, count);
101
+ }
102
+ static Term router_features(Env e, const float *x, const float *g, const float *b, u32 lo, u32 hi) {
103
+ if (hi - lo == 1) {
104
+ Loc node = heap_alloc(e, cls_fit(3));
105
+ e.mem[node] = f32_rewrap(x[lo]);
106
+ e.mem[node + 1] = f32_rewrap(g[lo]);
107
+ e.mem[node + 2] = f32_rewrap(b[lo]);
108
+ return term_ctr(CID_FEATURE, node);
109
+ }
110
+ u32 mid = lo + (hi - lo) / 2;
111
+ Term left = router_features(e, x, g, b, lo, mid);
112
+ Term right = router_features(e, x, g, b, mid, hi);
113
+ return io_node(e, CID_BRANCH, left, right);
114
+ }
115
+ static void router_unfeatures(Env e, Term tree, float *out, u32 lo, u32 hi) {
116
+ Loc node = term_loc(tree);
117
+ if (hi - lo == 1) {
118
+ out[lo] = f32_unbox(e.mem[node]);
119
+ heap_free(e, cls_fit(3), node);
120
+ } else {
121
+ u32 mid = lo + (hi - lo) / 2;
122
+ router_unfeatures(e, e.mem[node], out, lo, mid);
123
+ router_unfeatures(e, e.mem[node + 1], out, mid, hi);
124
+ heap_free(e, cls_fit(2), node);
125
+ }
126
+ }
127
+ static Term router_vector(Env e, const float *x, const float *g, const float *b, u32 width) {
128
+ Term tail = term_ctr(CID_VNIL, 0);
129
+ for (u32 i = width; i-- > 0;) {
130
+ Loc node = heap_alloc(e, cls_fit(4));
131
+ e.mem[node] = f32_rewrap(x[i]);
132
+ e.mem[node + 1] = f32_rewrap(g[i]);
133
+ e.mem[node + 2] = f32_rewrap(b[i]);
134
+ e.mem[node + 3] = tail;
135
+ tail = term_ctr(CID_VCONS, node);
136
+ }
137
+ return tail;
138
+ }
139
+ static Term router_rows(Env e, const float *x, const float *g, const float *b,
140
+ u32 width, u32 lo, u32 hi) {
141
+ if (hi - lo == 1) {
142
+ Loc node = heap_alloc(e, cls_fit(1));
143
+ e.mem[node] = router_vector(e, x + lo * width, g, b, width);
144
+ return term_ctr(CID_NLEAF, node);
145
+ }
146
+ u32 mid = lo + (hi - lo) / 2;
147
+ Term left = router_rows(e, x, g, b, width, lo, mid);
148
+ Term right = router_rows(e, x, g, b, width, mid, hi);
149
+ return io_node(e, CID_NFORK, left, right);
150
+ }
151
+ static void router_unrows(Env e, Term tree, float *out, u32 width, u32 lo, u32 hi) {
152
+ Loc node = term_loc(tree);
153
+ if (hi - lo == 1) {
154
+ Term vector = e.mem[node];
155
+ // The flat affine loop builds the result in reverse order.
156
+ for (u32 i = width; i-- > 0;) {
157
+ Loc item = term_loc(vector);
158
+ out[lo * width + i] = f32_unbox(e.mem[item]);
159
+ vector = e.mem[item + 3];
160
+ heap_free(e, cls_fit(4), item);
161
+ }
162
+ heap_free(e, cls_fit(1), node);
163
+ } else {
164
+ u32 mid = lo + (hi - lo) / 2;
165
+ router_unrows(e, e.mem[node], out, width, lo, mid);
166
+ router_unrows(e, e.mem[node + 1], out, width, mid, hi);
167
+ heap_free(e, cls_fit(2), node);
168
+ }
169
+ }
170
+ int julia_router_layernorm(const float *x, const float *gamma, const float *beta,
171
+ uint32_t rows, uint32_t width, float epsilon, float *out) {
172
+ if (!rows || !width || width > 65536 || rows > 1048576 / width || !out ||
173
+ !isfinite(epsilon) || epsilon <= 0) return -1;
174
+ int status = router_valid(x, rows * width);
175
+ if (!status) status = router_valid(gamma, width);
176
+ if (!status) status = router_valid(beta, width);
177
+ if (status) return status;
178
+ pthread_mutex_lock(&router_mutex);
179
+ Env e = router_env();
180
+ Term tree = router_rows(e, x, gamma, beta, width, 0, rows);
181
+ Loc task = task_node(e, FID_NORM_ROWS, TERM_HOLE, 0, 0);
182
+ e.mem[task] = tree;
183
+ e.mem[task + 1] = f32_rewrap((float)width);
184
+ e.mem[task + 2] = f32_rewrap(epsilon);
185
+ Term result = corpus_eval(e.mem, term_tsk(FID_NORM_ROWS, task));
186
+ router_unrows(e, result, out, width, 0, rows);
187
+ pthread_mutex_unlock(&router_mutex);
188
+ return router_valid(out, rows * width);
189
+ }
190
+
191
+ /* Resident immutable weights are built once and borrowed by every forward. */
192
+ typedef struct { Term tree; u32 inputs; u32 outputs; bool packed; } RouterMatrix;
193
+ static Term router_matrix(Env e, const float *w, u32 inputs, u32 outputs, u32 lo, u32 hi) {
194
+ if (hi - lo == 1) {
195
+ Term tail = term_ctr(CID_JWEND, 0);
196
+ for (u32 j = inputs; j-- > 0;) {
197
+ Loc node = heap_alloc(e, cls_fit(8));
198
+ for (u32 k = 0; k < 7; ++k)
199
+ e.mem[node + k] = f32_rewrap(lo * 7 + k < outputs ? w[(lo * 7 + k) * inputs + j] : 0.f);
200
+ e.mem[node + 7] = tail;
201
+ tail = term_ctr(CID_JWVAL, node);
202
+ }
203
+ Loc leaf = heap_alloc(e, cls_fit(1));
204
+ e.mem[leaf] = tail;
205
+ return term_ctr(CID_JMLEAF, leaf);
206
+ }
207
+ u32 mid = lo + (hi - lo) / 2;
208
+ Term left = router_matrix(e, w, inputs, outputs, lo, mid);
209
+ Term right = router_matrix(e, w, inputs, outputs, mid, hi);
210
+ return io_node(e, CID_JMFORK, left, right);
211
+ }
212
+ static void router_free_matrix(Env e, Term tree) {
213
+ Loc node = term_loc(tree);
214
+ if (term_aux(tree) == CID_JMLEAF) {
215
+ Term ws = e.mem[node];
216
+ while (term_aux(ws) == CID_JWVAL) {
217
+ Loc item = term_loc(ws);
218
+ ws = e.mem[item + 7];
219
+ heap_free(e, cls_fit(8), item);
220
+ }
221
+ heap_free(e, cls_fit(1), node);
222
+ } else {
223
+ router_free_matrix(e, e.mem[node]);
224
+ router_free_matrix(e, e.mem[node + 1]);
225
+ heap_free(e, cls_fit(2), node);
226
+ }
227
+ }
228
+ /* Packed FP32 buffers: padding is explicit because Bend array indices wrap. */
229
+ static Term router_buffer_raw(Env e, u32 cells) {
230
+ Cls depth = cls_fit(cells);
231
+ Loc loc = heap_alloc(e, buf_wcls(depth));
232
+ return term_buf(depth, loc);
233
+ }
234
+ static Term router_packed_weights(Env e, const float *w, u32 inputs, u32 outputs) {
235
+ u32 cols = (outputs + 7) / 8 * 8;
236
+ Term array = router_buffer_raw(e, inputs * cols);
237
+ memset((void *)(e.mem + term_loc(array)), 0, (size_t)inputs * cols * sizeof(float));
238
+ float *data = (float *)(e.mem + term_loc(array));
239
+ for (u32 j = 0; j < inputs; ++j)
240
+ for (u32 col = 0; col < outputs; ++col)
241
+ data[(col / 8 * inputs + j) * 8 + col % 8] = w[col * inputs + j];
242
+ return array;
243
+ }
244
+ /* Layout conversion only: transpose eight FP32 rows without doing arithmetic. */
245
+ static void router_pack_input(float *restrict dst, const float *restrict src, u32 rows, u32 width) {
246
+ u32 row = 0;
247
+ #if defined(__AVX2__)
248
+ for (; row + 8 <= rows; row += 8) {
249
+ u32 j = 0;
250
+ for (; j + 8 <= width; j += 8) {
251
+ __m256 r[8], t[8], v[8];
252
+ for (u32 k = 0; k < 8; ++k) r[k] = _mm256_loadu_ps(src + (row + k) * width + j);
253
+ for (u32 k = 0; k < 8; k += 2) {
254
+ t[k] = _mm256_unpacklo_ps(r[k], r[k+1]);
255
+ t[k+1] = _mm256_unpackhi_ps(r[k], r[k+1]);
256
+ }
257
+ for (u32 k = 0; k < 8; k += 4) {
258
+ v[k] = _mm256_shuffle_ps(t[k], t[k+2], 0x44);
259
+ v[k+1] = _mm256_shuffle_ps(t[k], t[k+2], 0xee);
260
+ v[k+2] = _mm256_shuffle_ps(t[k+1], t[k+3], 0x44);
261
+ v[k+3] = _mm256_shuffle_ps(t[k+1], t[k+3], 0xee);
262
+ }
263
+ for (u32 k = 0; k < 4; ++k) {
264
+ _mm256_storeu_ps(dst + (row / 8 * width + j + k) * 8,
265
+ _mm256_permute2f128_ps(v[k], v[k+4], 0x20));
266
+ _mm256_storeu_ps(dst + (row / 8 * width + j + k + 4) * 8,
267
+ _mm256_permute2f128_ps(v[k], v[k+4], 0x31));
268
+ }
269
+ }
270
+ for (; j < width; ++j)
271
+ for (u32 k = 0; k < 8; ++k) dst[(row / 8 * width + j) * 8 + k] = src[(row+k) * width + j];
272
+ }
273
+ #endif
274
+ for (; row < rows; ++row)
275
+ for (u32 j = 0; j < width; ++j)
276
+ dst[(row / 8 * width + j) * 8 + row % 8] = src[row * width + j];
277
+ }
278
+ static void router_packed_linear(Env e, RouterMatrix *matrix, const float *x, u32 rows, float *out) {
279
+ double begin = router_profile_enabled ? router_clock_ms() : 0.0;
280
+ u32 cols = (matrix->outputs + 7) / 8 * 8;
281
+ u32 padded_rows = (rows + 7) / 8 * 8;
282
+ u32 count = padded_rows / 8 * (cols / 8);
283
+ Term input = router_buffer_raw(e, padded_rows * matrix->inputs);
284
+ Term output = router_buffer_raw(e, padded_rows * cols);
285
+ float *input_data = (float *)(e.mem + term_loc(input));
286
+ if (rows % 8)
287
+ memset(input_data + (rows / 8) * matrix->inputs * 8, 0, matrix->inputs * 8 * sizeof(float));
288
+ router_pack_input(input_data, x, rows, matrix->inputs);
289
+ Loc task = task_node(e, FID_PACKED_GEMM, TERM_HOLE, 0, 0);
290
+ e.mem[task + 0] = input;
291
+ e.mem[task + 1] = matrix->tree;
292
+ e.mem[task + 2] = output;
293
+ e.mem[task + 3] = 0;
294
+ e.mem[task + 4] = count;
295
+ e.mem[task + 5] = cols / 8;
296
+ const char *grain_setting = getenv("JULIA_BEND_TILE_GRAIN");
297
+ u32 automatic_grain = (count + pool_size * 2 - 1) / (pool_size * 2);
298
+ u32 grain = grain_setting ? strtoul(grain_setting, NULL, 10) : automatic_grain;
299
+ if (!grain || grain > 65536) grain = automatic_grain;
300
+ e.mem[task + 6] = grain;
301
+ e.mem[task + 7] = matrix->inputs;
302
+ e.mem[task + 8] = (count <= grain);
303
+ /* PackedBuffers is flattened to three result words by Bend 2.0.27.
304
+ * corpus_eval clears the ready flag but leaves these words intact. The
305
+ * runtime mutex prevents another evaluation from overwriting the mailbox. */
306
+ double dispatched = router_profile_enabled ? router_clock_ms() : 0.0;
307
+ corpus_eval(e.mem, term_tsk(FID_PACKED_GEMM, task));
308
+ double evaluated = router_profile_enabled ? router_clock_ms() : 0.0;
309
+ input = e.mem[H_ROOT_WORD];
310
+ matrix->tree = e.mem[H_ROOT_WORD + 1];
311
+ output = e.mem[H_ROOT_WORD + 2];
312
+ const float *data = (const float *)(e.mem + blk_loc(e.mem, output));
313
+ for (u32 row = 0; row < rows; ++row) {
314
+ const float *tiles = data + (row / 8 * (cols / 8)) * 64 + row % 8 * 8;
315
+ u32 col = 0;
316
+ for (; col + 8 <= matrix->outputs; col += 8)
317
+ memcpy(out + row * matrix->outputs + col, tiles + (col / 8) * 64, 8 * sizeof(float));
318
+ if (col < matrix->outputs)
319
+ memcpy(out + row * matrix->outputs + col, tiles + (col / 8) * 64,
320
+ (matrix->outputs - col) * sizeof(float));
321
+ }
322
+ blk_free(e, input);
323
+ blk_free(e, output);
324
+ if (router_profile_enabled) {
325
+ router_profile_totals[0] += dispatched - begin;
326
+ router_profile_totals[1] += evaluated - dispatched;
327
+ router_profile_totals[2] += router_clock_ms() - evaluated;
328
+ router_profile_totals[3] += 1;
329
+ }
330
+ }
331
+ void *julia_router_matrix_create(const float *weights, u32 inputs, u32 outputs) {
332
+ if (!inputs || !outputs || inputs > 65536 || outputs > 65536 ||
333
+ (u64)inputs * outputs > 16777216) return NULL;
334
+ if (!weights) return NULL;
335
+ for (u64 i = 0; i < (u64)inputs * outputs; ++i) if (!isfinite(weights[i])) return NULL;
336
+ RouterMatrix *matrix = malloc(sizeof(*matrix));
337
+ if (!matrix) return NULL;
338
+ pthread_mutex_lock(&router_mutex);
339
+ Env e = router_env();
340
+ matrix->inputs = inputs;
341
+ matrix->outputs = outputs;
342
+ const char *kernel = getenv("JULIA_BEND_KERNEL");
343
+ matrix->packed = !kernel || strcmp(kernel, "tree") != 0;
344
+ matrix->tree = matrix->packed ? router_packed_weights(e, weights, inputs, outputs)
345
+ : router_matrix(e, weights, inputs, outputs, 0, (outputs + 6) / 7);
346
+ pthread_mutex_unlock(&router_mutex);
347
+ return matrix;
348
+ }
349
+ void julia_router_matrix_free(void *handle) {
350
+ if (!handle) return;
351
+ RouterMatrix *matrix = handle;
352
+ pthread_mutex_lock(&router_mutex);
353
+ if (matrix->packed) blk_free(router_env(), matrix->tree);
354
+ else router_free_matrix(router_env(), matrix->tree);
355
+ pthread_mutex_unlock(&router_mutex);
356
+ free(matrix);
357
+ }
358
+ static void router_tile_result(Env e, Term tree, float *out, u32 rows, u32 outputs, u32 lo, u32 hi) {
359
+ Loc node = term_loc(tree);
360
+ if (hi - lo == 1) {
361
+ for (u32 row = 0; row < rows; ++row)
362
+ for (u32 k = 0; k < 7 && lo * 7 + k < outputs; ++k)
363
+ out[row * outputs + lo * 7 + k] = f32_unbox(e.mem[node + row * 7 + k]);
364
+ heap_free(e, cls_fit(28), node);
365
+ } else {
366
+ u32 mid = lo + (hi - lo) / 2;
367
+ router_tile_result(e, e.mem[node], out, rows, outputs, lo, mid);
368
+ router_tile_result(e, e.mem[node + 1], out, rows, outputs, mid, hi);
369
+ heap_free(e, cls_fit(2), node);
370
+ }
371
+ }
372
+ static Term router_input_tiles(Env e, const float *x, u32 rows, u32 inputs, u32 lo, u32 hi) {
373
+ if (hi - lo == 1) {
374
+ Term xs = term_ctr(CID_JTEND, 0);
375
+ for (u32 i = inputs; i-- > 0;) {
376
+ Loc item = heap_alloc(e, cls_fit(5));
377
+ for (u32 k = 0; k < 4; ++k)
378
+ e.mem[item + k] = f32_rewrap(lo * 4 + k < rows ? x[(lo * 4 + k) * inputs + i] : 0.f);
379
+ e.mem[item + 4] = xs;
380
+ xs = term_ctr(CID_JTVAL, item);
381
+ }
382
+ Loc leaf = heap_alloc(e, cls_fit(1));
383
+ e.mem[leaf] = xs;
384
+ return term_ctr(CID_JILEAF, leaf);
385
+ }
386
+ u32 mid = lo + (hi - lo) / 2;
387
+ Term left = router_input_tiles(e, x, rows, inputs, lo, mid);
388
+ Term right = router_input_tiles(e, x, rows, inputs, mid, hi);
389
+ return io_node(e, CID_JIFORK, left, right);
390
+ }
391
+ static void router_free_tiles(Env e, Term tree) {
392
+ Loc node = term_loc(tree);
393
+ if (term_aux(tree) == CID_JILEAF) {
394
+ Term xs = e.mem[node];
395
+ while (term_aux(xs) == CID_JTVAL) {
396
+ Loc item = term_loc(xs);
397
+ xs = e.mem[item + 4];
398
+ heap_free(e, cls_fit(5), item);
399
+ }
400
+ heap_free(e, cls_fit(1), node);
401
+ } else {
402
+ router_free_tiles(e, e.mem[node]);
403
+ router_free_tiles(e, e.mem[node + 1]);
404
+ heap_free(e, cls_fit(2), node);
405
+ }
406
+ }
407
+ static void router_output_tiles(Env e, Term tree, float *out, u32 rows, u32 outputs, u32 lo, u32 hi) {
408
+ Loc node = term_loc(tree);
409
+ if (hi - lo == 1) {
410
+ u32 count = rows - lo * 4 < 4 ? rows - lo * 4 : 4;
411
+ router_tile_result(e, e.mem[node], out + lo * 4 * outputs, count, outputs, 0, (outputs + 6) / 7);
412
+ heap_free(e, cls_fit(1), node);
413
+ } else {
414
+ u32 mid = lo + (hi - lo) / 2;
415
+ router_output_tiles(e, e.mem[node], out, rows, outputs, lo, mid);
416
+ router_output_tiles(e, e.mem[node + 1], out, rows, outputs, mid, hi);
417
+ heap_free(e, cls_fit(2), node);
418
+ }
419
+ }
420
+ static int router_linear(void *handle, const float *x, u32 rows, float *out, bool validate) {
421
+ if (!handle || !x || !out || !rows) return -1;
422
+ RouterMatrix *matrix = handle;
423
+ if (rows > 1048576 / matrix->inputs || rows > 1048576 / matrix->outputs) return -1;
424
+ int status = validate ? router_valid(x, rows * matrix->inputs) : 0;
425
+ if (status) return status;
426
+ pthread_mutex_lock(&router_mutex);
427
+ Env e = router_env();
428
+ if (matrix->packed) {
429
+ router_packed_linear(e, matrix, x, rows, out);
430
+ pthread_mutex_unlock(&router_mutex);
431
+ return validate ? router_valid(out, rows * matrix->outputs) : 0;
432
+ }
433
+ u32 count = (rows + 3) / 4;
434
+ Term xs = router_input_tiles(e, x, rows, matrix->inputs, 0, count);
435
+ Loc task = task_node(e, FID_DENSE_BATCH, TERM_HOLE, 0, 0);
436
+ e.mem[task] = xs;
437
+ e.mem[task + 1] = matrix->tree;
438
+ Term result = corpus_eval(e.mem, term_tsk(FID_DENSE_BATCH, task));
439
+ router_output_tiles(e, result, out, rows, matrix->outputs, 0, count);
440
+ router_free_tiles(e, xs);
441
+ pthread_mutex_unlock(&router_mutex);
442
+ return validate ? router_valid(out, rows * matrix->outputs) : 0;
443
+ }
444
+
445
+ int julia_router_linear(void *handle, const float *x, u32 rows, float *out) {
446
+ return router_linear(handle, x, rows, out, true);
447
+ }
448
+ /* Internal engine path: weights are validated once, logits at the model boundary.
449
+ * Shape, pointer, lifetime and buffer bounds checks remain mandatory. */
450
+ int julia_router_linear_inference(void *handle, const float *x, u32 rows, float *out) {
451
+ return router_linear(handle, x, rows, out, false);
452
+ }
julia/router/native/router.bend ADDED
@@ -0,0 +1,360 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import Base
2
+
3
+ type Candidates is Data:
4
+ Leaf{index: U32, score: F32}
5
+ Fork{left: Candidates, right: Candidates}
6
+
7
+ def pick(flag: Bool, a: Candidates, b: Candidates) -> Candidates:
8
+ match flag:
9
+ case True{}:
10
+ a
11
+ case False{}:
12
+ b
13
+
14
+ def better(a: Candidates, b: Candidates) -> Candidates:
15
+ match a b:
16
+ case Leaf{+i, +x} Leaf{+j, +y}:
17
+ pick(F32.is_gt(x, y) || (F32.is_eq(x, y) && U32.is_le(i, j)), Leaf{i, x}, Leaf{j, y})
18
+ case _ _:
19
+ a
20
+
21
+ def winner(tree: Candidates) -> Candidates:
22
+ match tree:
23
+ case Leaf{i, x}:
24
+ Leaf{i, x}
25
+ case Fork{left, right}:
26
+ a b = winner(left) winner(right)
27
+ better(a, b)
28
+
29
+ def input() -> IO(Candidates):
30
+ import "./bridge.c"
31
+
32
+ def output(tree: Candidates) -> IO(Unit):
33
+ import "./bridge.c"
34
+
35
+ # Transformer math uses the same affine trees. Forks expose independent work.
36
+ def exponentials(tree: Candidates, +maximum: F32) -> Candidates:
37
+ match tree:
38
+ case Leaf{i, x}:
39
+ Leaf{i, F32.exp(F32.sub(x, maximum))}
40
+ case Fork{left, right}:
41
+ a b = exponentials(left, maximum) exponentials(right, maximum)
42
+ Fork{a, b}
43
+
44
+ def total(tree: Candidates) -> F32:
45
+ match tree:
46
+ case Leaf{i, x}:
47
+ x
48
+ case Fork{left, right}:
49
+ a b = total(left) total(right)
50
+ F32.add(a, b)
51
+
52
+ def divide(tree: Candidates, +sum: F32) -> Candidates:
53
+ match tree:
54
+ case Leaf{i, x}:
55
+ Leaf{i, F32.div(x, sum)}
56
+ case Fork{left, right}:
57
+ a b = divide(left, sum) divide(right, sum)
58
+ Fork{a, b}
59
+
60
+ def softmax_best(tree: Candidates, best: Candidates) -> Candidates:
61
+ match best:
62
+ case Leaf{i, maximum}:
63
+ +values = exponentials(tree, maximum)
64
+ divide(values, total(values))
65
+ case Fork{left, right}:
66
+ tree
67
+
68
+ def softmax(+tree: Candidates) -> Candidates:
69
+ softmax_best(tree, winner(tree))
70
+
71
+ type Features is Data:
72
+ Feature{x: F32, gamma: F32, beta: F32}
73
+ Branch{left: Features, right: Features}
74
+
75
+ def feature_sum(tree: Features) -> F32:
76
+ match tree:
77
+ case Feature{x, gamma, beta}:
78
+ x
79
+ case Branch{left, right}:
80
+ a b = feature_sum(left) feature_sum(right)
81
+ F32.add(a, b)
82
+
83
+ def variance(tree: Features, +mean: F32) -> F32:
84
+ match tree:
85
+ case Feature{x, gamma, beta}:
86
+ F32.square(F32.sub(x, mean))
87
+ case Branch{left, right}:
88
+ a b = variance(left, mean) variance(right, mean)
89
+ F32.add(a, b)
90
+
91
+ def affine(tree: Features, +mean: F32, +scale: F32) -> Features:
92
+ match tree:
93
+ case Feature{x, +gamma, +beta}:
94
+ Feature{F32.add(F32.mul(F32.mul(F32.sub(x, mean), scale), gamma), beta), gamma, beta}
95
+ case Branch{left, right}:
96
+ a b = affine(left, mean, scale) affine(right, mean, scale)
97
+ Branch{a, b}
98
+
99
+ def layernorm(+tree: Features, +count: F32, epsilon: F32) -> Features:
100
+ +mean = F32.div(feature_sum(tree), count)
101
+ scale = F32.div(1.0, F32.sqrt(F32.add(F32.div(variance(tree, mean), count), epsilon)))
102
+ affine(tree, mean, scale)
103
+
104
+ def features() -> IO(Features):
105
+ import "./bridge.c"
106
+
107
+ def write_features(tree: Features) -> IO(Unit):
108
+ import "./bridge.c"
109
+
110
+ # Coarse parallelism: fork over independent rows, flat tail loops within a row.
111
+ # Scalar-level tasks cost more to schedule than the arithmetic they perform.
112
+ type Vector is Data:
113
+ VNil{}
114
+ VCons{x: F32, gamma: F32, beta: F32, next: Vector}
115
+
116
+ type NormRows is Data:
117
+ NLeaf{values: Vector}
118
+ NFork{left: NormRows, right: NormRows}
119
+
120
+ def vector_sum(xs: Vector, acc: F32) -> F32:
121
+ match xs:
122
+ case VNil{}:
123
+ acc
124
+ case VCons{x, g, b, next}:
125
+ vector_sum(next, F32.add(acc, x))
126
+
127
+ def vector_variance(xs: Vector, +mean: F32, acc: F32) -> F32:
128
+ match xs:
129
+ case VNil{}:
130
+ acc
131
+ case VCons{x, g, b, next}:
132
+ vector_variance(next, mean, F32.add(acc, F32.square(F32.sub(x, mean))))
133
+
134
+ def vector_affine(xs: Vector, +mean: F32, +scale: F32, out: Vector) -> Vector:
135
+ match xs:
136
+ case VNil{}:
137
+ out
138
+ case VCons{x, +g, +b, next}:
139
+ value = F32.add(F32.mul(F32.mul(F32.sub(x, mean), scale), g), b)
140
+ vector_affine(next, mean, scale, VCons{value, g, b, out})
141
+
142
+ def norm_vector(+xs: Vector, +count: F32, epsilon: F32) -> Vector:
143
+ +mean = F32.div(vector_sum(xs, 0.0), count)
144
+ scale = F32.div(1.0, F32.sqrt(F32.add(F32.div(vector_variance(xs, mean, 0.0), count), epsilon)))
145
+ vector_affine(xs, mean, scale, VNil{})
146
+
147
+ def norm_rows(xs: NormRows, +count: F32, +epsilon: F32) -> NormRows:
148
+ match xs:
149
+ case NLeaf{values}:
150
+ NLeaf{norm_vector(values, count, epsilon)}
151
+ case NFork{left, right}:
152
+ a b = norm_rows(left, count, epsilon) norm_rows(right, count, epsilon)
153
+ NFork{a, b}
154
+
155
+ def batch_features() -> IO(NormRows):
156
+ import "./bridge.c"
157
+
158
+ def write_batch(tree: NormRows) -> IO(Unit):
159
+ import "./bridge.c"
160
+
161
+ # Resident dense projections. Each leaf computes four output channels in one
162
+ # flat loop; only independent output tiles fork. Inputs and weights are borrowed.
163
+ type DenseWeights is Data:
164
+ JWEnd{}
165
+ JWVal{w0: F32, w1: F32, w2: F32, w3: F32, w4: F32, w5: F32, w6: F32, next: DenseWeights}
166
+
167
+ type JuliaMatrix is Data:
168
+ JMLeaf{weights: DenseWeights}
169
+ JMFork{left: JuliaMatrix, right: JuliaMatrix}
170
+
171
+ def dense_weights() -> IO(JuliaMatrix):
172
+ import "./bridge.c"
173
+
174
+ # A 4x7 tile fits seven weights plus a link in one eight-word heap allocation.
175
+ # Each weight is reused across four tokens; no half-empty 16-word allocations.
176
+ type TileInput is Data:
177
+ JTEnd{}
178
+ JTVal{a: F32, b: F32, c: F32, d: F32, next: TileInput}
179
+
180
+ type TileResult is Data:
181
+ JTLeaf{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32, v8: F32, v9: F32, v10: F32, v11: F32, v12: F32, v13: F32, v14: F32, v15: F32, v16: F32, v17: F32, v18: F32, v19: F32, v20: F32, v21: F32, v22: F32, v23: F32, v24: F32, v25: F32, v26: F32, v27: F32}
182
+ JTFork{left: TileResult, right: TileResult}
183
+
184
+ def dot_tile(xs: TileInput, ws: DenseWeights, s0: F32, s1: F32, s2: F32, s3: F32, s4: F32, s5: F32, s6: F32, s7: F32, s8: F32, s9: F32, s10: F32, s11: F32, s12: F32, s13: F32, s14: F32, s15: F32, s16: F32, s17: F32, s18: F32, s19: F32, s20: F32, s21: F32, s22: F32, s23: F32, s24: F32, s25: F32, s26: F32, s27: F32) -> TileResult:
185
+ match xs ws:
186
+ case JTVal{+x0, +x1, +x2, +x3, xt} JWVal{+w0, +w1, +w2, +w3, +w4, +w5, +w6, wt}:
187
+ dot_tile(xt, wt, F32.add(s0, F32.mul(x0, w0)),
188
+ F32.add(s1, F32.mul(x0, w1)),
189
+ F32.add(s2, F32.mul(x0, w2)),
190
+ F32.add(s3, F32.mul(x0, w3)),
191
+ F32.add(s4, F32.mul(x0, w4)),
192
+ F32.add(s5, F32.mul(x0, w5)),
193
+ F32.add(s6, F32.mul(x0, w6)),
194
+ F32.add(s7, F32.mul(x1, w0)),
195
+ F32.add(s8, F32.mul(x1, w1)),
196
+ F32.add(s9, F32.mul(x1, w2)),
197
+ F32.add(s10, F32.mul(x1, w3)),
198
+ F32.add(s11, F32.mul(x1, w4)),
199
+ F32.add(s12, F32.mul(x1, w5)),
200
+ F32.add(s13, F32.mul(x1, w6)),
201
+ F32.add(s14, F32.mul(x2, w0)),
202
+ F32.add(s15, F32.mul(x2, w1)),
203
+ F32.add(s16, F32.mul(x2, w2)),
204
+ F32.add(s17, F32.mul(x2, w3)),
205
+ F32.add(s18, F32.mul(x2, w4)),
206
+ F32.add(s19, F32.mul(x2, w5)),
207
+ F32.add(s20, F32.mul(x2, w6)),
208
+ F32.add(s21, F32.mul(x3, w0)),
209
+ F32.add(s22, F32.mul(x3, w1)),
210
+ F32.add(s23, F32.mul(x3, w2)),
211
+ F32.add(s24, F32.mul(x3, w3)),
212
+ F32.add(s25, F32.mul(x3, w4)),
213
+ F32.add(s26, F32.mul(x3, w5)),
214
+ F32.add(s27, F32.mul(x3, w6)))
215
+ case _ _:
216
+ JTLeaf{s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27}
217
+
218
+ def dense_tile(+xs: TileInput, weights: JuliaMatrix) -> TileResult:
219
+ match weights:
220
+ case JMLeaf{ws}:
221
+ dot_tile(xs, ws, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
222
+ case JMFork{left, right}:
223
+ a b = dense_tile(xs, left) dense_tile(xs, right)
224
+ JTFork{a, b}
225
+
226
+ def tile_input() -> IO(TileInput):
227
+ import "./bridge.c"
228
+
229
+ def tile_output(result: TileResult) -> IO(Unit):
230
+ import "./bridge.c"
231
+
232
+ # One fork tree spans token tiles and output tiles; one runtime dispatch per matrix.
233
+ type InputTiles is Data:
234
+ JILeaf{tile: TileInput}
235
+ JIFork{left: InputTiles, right: InputTiles}
236
+
237
+ type OutputTiles is Data:
238
+ JOLeaf{tile: TileResult}
239
+ JOFork{left: OutputTiles, right: OutputTiles}
240
+
241
+ def dense_batch(tiles: InputTiles, +weights: JuliaMatrix) -> OutputTiles:
242
+ match tiles:
243
+ case JILeaf{tile}:
244
+ JOLeaf{dense_tile(tile, weights)}
245
+ case JIFork{left, right}:
246
+ a b = dense_batch(left, weights) dense_batch(right, weights)
247
+ JOFork{a, b}
248
+
249
+ def input_tiles() -> IO(InputTiles):
250
+ import "./bridge.c"
251
+
252
+ def output_tiles(result: OutputTiles) -> IO(Unit):
253
+ import "./bridge.c"
254
+
255
+ # Packed tiles expose contiguous lanes to the native compiler. All arithmetic
256
+ # stays in Bend. Shared inputs are read-only; output tiles are disjoint.
257
+ # Unsafe loops terminate at the positive, validated dimensions in the bridge.
258
+ type Input8 is Data:
259
+ JInput8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32}
260
+
261
+ type Weight8 is Data:
262
+ JWeight8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32}
263
+
264
+ type Output64 is Data:
265
+ JOutput64{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32, v8: F32, v9: F32, v10: F32, v11: F32, v12: F32, v13: F32, v14: F32, v15: F32, v16: F32, v17: F32, v18: F32, v19: F32, v20: F32, v21: F32, v22: F32, v23: F32, v24: F32, v25: F32, v26: F32, v27: F32, v28: F32, v29: F32, v30: F32, v31: F32, v32: F32, v33: F32, v34: F32, v35: F32, v36: F32, v37: F32, v38: F32, v39: F32, v40: F32, v41: F32, v42: F32, v43: F32, v44: F32, v45: F32, v46: F32, v47: F32, v48: F32, v49: F32, v50: F32, v51: F32, v52: F32, v53: F32, v54: F32, v55: F32, v56: F32, v57: F32, v58: F32, v59: F32, v60: F32, v61: F32, v62: F32, v63: F32}
266
+
267
+ type PackedBuffers is Type:
268
+ JBuffers{x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>}
269
+
270
+ @unsafe
271
+ def packed_dot(x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>, +row: U32, +col: U32, +cols: U32, +inner: U32, +end: U32, +j: U32, +s0: F32, +s1: F32, +s2: F32, +s3: F32, +s4: F32, +s5: F32, +s6: F32, +s7: F32, +s8: F32, +s9: F32, +s10: F32, +s11: F32, +s12: F32, +s13: F32, +s14: F32, +s15: F32, +s16: F32, +s17: F32, +s18: F32, +s19: F32, +s20: F32, +s21: F32, +s22: F32, +s23: F32, +s24: F32, +s25: F32, +s26: F32, +s27: F32, +s28: F32, +s29: F32, +s30: F32, +s31: F32, +s32: F32, +s33: F32, +s34: F32, +s35: F32, +s36: F32, +s37: F32, +s38: F32, +s39: F32, +s40: F32, +s41: F32, +s42: F32, +s43: F32, +s44: F32, +s45: F32, +s46: F32, +s47: F32, +s48: F32, +s49: F32, +s50: F32, +s51: F32, +s52: F32, +s53: F32, +s54: F32, +s55: F32, +s56: F32, +s57: F32, +s58: F32, +s59: F32, +s60: F32, +s61: F32, +s62: F32, +s63: F32, active: Bool) -> PackedBuffers:
272
+ match active:
273
+ case True{}:
274
+ packed_read_x(Array.get(Input8, x, U32.add(U32.mul(row, inner), j)), w, o, row, col, cols, inner, end, j, s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27, s28, s29, s30, s31, s32, s33, s34, s35, s36, s37, s38, s39, s40, s41, s42, s43, s44, s45, s46, s47, s48, s49, s50, s51, s52, s53, s54, s55, s56, s57, s58, s59, s60, s61, s62, s63)
275
+ case False{}:
276
+ o = Array.set(Output64, o, U32.add(U32.mul(row, cols), col), JOutput64{s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27, s28, s29, s30, s31, s32, s33, s34, s35, s36, s37, s38, s39, s40, s41, s42, s43, s44, s45, s46, s47, s48, s49, s50, s51, s52, s53, s54, s55, s56, s57, s58, s59, s60, s61, s62, s63})
277
+ packed_step(x, w, o, U32.add(U32.add(U32.mul(row, cols), col), 1), end, cols, inner, U32.is_lt(U32.add(U32.add(U32.mul(row, cols), col), 1), end))
278
+
279
+ @unsafe
280
+ def packed_read_x(pair: Array<Input8> & Input8, w: Array<Weight8>, o: Array<Output64>, +row: U32, +col: U32, +cols: U32, +inner: U32, +end: U32, +j: U32, +s0: F32, +s1: F32, +s2: F32, +s3: F32, +s4: F32, +s5: F32, +s6: F32, +s7: F32, +s8: F32, +s9: F32, +s10: F32, +s11: F32, +s12: F32, +s13: F32, +s14: F32, +s15: F32, +s16: F32, +s17: F32, +s18: F32, +s19: F32, +s20: F32, +s21: F32, +s22: F32, +s23: F32, +s24: F32, +s25: F32, +s26: F32, +s27: F32, +s28: F32, +s29: F32, +s30: F32, +s31: F32, +s32: F32, +s33: F32, +s34: F32, +s35: F32, +s36: F32, +s37: F32, +s38: F32, +s39: F32, +s40: F32, +s41: F32, +s42: F32, +s43: F32, +s44: F32, +s45: F32, +s46: F32, +s47: F32, +s48: F32, +s49: F32, +s50: F32, +s51: F32, +s52: F32, +s53: F32, +s54: F32, +s55: F32, +s56: F32, +s57: F32, +s58: F32, +s59: F32, +s60: F32, +s61: F32, +s62: F32, +s63: F32) -> PackedBuffers:
281
+ (x, input) = pair
282
+ JInput8{+x0, +x1, +x2, +x3, +x4, +x5, +x6, +x7} = input
283
+ packed_read_w(Array.get(Weight8, w, U32.add(U32.mul(col, inner), j)), x, o, row, col, cols, inner, end, j, s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27, s28, s29, s30, s31, s32, s33, s34, s35, s36, s37, s38, s39, s40, s41, s42, s43, s44, s45, s46, s47, s48, s49, s50, s51, s52, s53, s54, s55, s56, s57, s58, s59, s60, s61, s62, s63, x0, x1, x2, x3, x4, x5, x6, x7)
284
+
285
+ @unsafe
286
+ def packed_read_w(pair: Array<Weight8> & Weight8, x: Array<Input8>, o: Array<Output64>, +row: U32, +col: U32, +cols: U32, +inner: U32, +end: U32, +j: U32, +s0: F32, +s1: F32, +s2: F32, +s3: F32, +s4: F32, +s5: F32, +s6: F32, +s7: F32, +s8: F32, +s9: F32, +s10: F32, +s11: F32, +s12: F32, +s13: F32, +s14: F32, +s15: F32, +s16: F32, +s17: F32, +s18: F32, +s19: F32, +s20: F32, +s21: F32, +s22: F32, +s23: F32, +s24: F32, +s25: F32, +s26: F32, +s27: F32, +s28: F32, +s29: F32, +s30: F32, +s31: F32, +s32: F32, +s33: F32, +s34: F32, +s35: F32, +s36: F32, +s37: F32, +s38: F32, +s39: F32, +s40: F32, +s41: F32, +s42: F32, +s43: F32, +s44: F32, +s45: F32, +s46: F32, +s47: F32, +s48: F32, +s49: F32, +s50: F32, +s51: F32, +s52: F32, +s53: F32, +s54: F32, +s55: F32, +s56: F32, +s57: F32, +s58: F32, +s59: F32, +s60: F32, +s61: F32, +s62: F32, +s63: F32, +x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, +x5: F32, +x6: F32, +x7: F32) -> PackedBuffers:
287
+ (w, weight) = pair
288
+ JWeight8{+w0, +w1, +w2, +w3, +w4, +w5, +w6, +w7} = weight
289
+ packed_dot(x, w, o, row, col, cols, inner, end, U32.add(j, 1), F32.add(s0, F32.mul(x0, w0)), F32.add(s1, F32.mul(x0, w1)), F32.add(s2, F32.mul(x0, w2)), F32.add(s3, F32.mul(x0, w3)), F32.add(s4, F32.mul(x0, w4)), F32.add(s5, F32.mul(x0, w5)), F32.add(s6, F32.mul(x0, w6)), F32.add(s7, F32.mul(x0, w7)), F32.add(s8, F32.mul(x1, w0)), F32.add(s9, F32.mul(x1, w1)), F32.add(s10, F32.mul(x1, w2)), F32.add(s11, F32.mul(x1, w3)), F32.add(s12, F32.mul(x1, w4)), F32.add(s13, F32.mul(x1, w5)), F32.add(s14, F32.mul(x1, w6)), F32.add(s15, F32.mul(x1, w7)), F32.add(s16, F32.mul(x2, w0)), F32.add(s17, F32.mul(x2, w1)), F32.add(s18, F32.mul(x2, w2)), F32.add(s19, F32.mul(x2, w3)), F32.add(s20, F32.mul(x2, w4)), F32.add(s21, F32.mul(x2, w5)), F32.add(s22, F32.mul(x2, w6)), F32.add(s23, F32.mul(x2, w7)), F32.add(s24, F32.mul(x3, w0)), F32.add(s25, F32.mul(x3, w1)), F32.add(s26, F32.mul(x3, w2)), F32.add(s27, F32.mul(x3, w3)), F32.add(s28, F32.mul(x3, w4)), F32.add(s29, F32.mul(x3, w5)), F32.add(s30, F32.mul(x3, w6)), F32.add(s31, F32.mul(x3, w7)), F32.add(s32, F32.mul(x4, w0)), F32.add(s33, F32.mul(x4, w1)), F32.add(s34, F32.mul(x4, w2)), F32.add(s35, F32.mul(x4, w3)), F32.add(s36, F32.mul(x4, w4)), F32.add(s37, F32.mul(x4, w5)), F32.add(s38, F32.mul(x4, w6)), F32.add(s39, F32.mul(x4, w7)), F32.add(s40, F32.mul(x5, w0)), F32.add(s41, F32.mul(x5, w1)), F32.add(s42, F32.mul(x5, w2)), F32.add(s43, F32.mul(x5, w3)), F32.add(s44, F32.mul(x5, w4)), F32.add(s45, F32.mul(x5, w5)), F32.add(s46, F32.mul(x5, w6)), F32.add(s47, F32.mul(x5, w7)), F32.add(s48, F32.mul(x6, w0)), F32.add(s49, F32.mul(x6, w1)), F32.add(s50, F32.mul(x6, w2)), F32.add(s51, F32.mul(x6, w3)), F32.add(s52, F32.mul(x6, w4)), F32.add(s53, F32.mul(x6, w5)), F32.add(s54, F32.mul(x6, w6)), F32.add(s55, F32.mul(x6, w7)), F32.add(s56, F32.mul(x7, w0)), F32.add(s57, F32.mul(x7, w1)), F32.add(s58, F32.mul(x7, w2)), F32.add(s59, F32.mul(x7, w3)), F32.add(s60, F32.mul(x7, w4)), F32.add(s61, F32.mul(x7, w5)), F32.add(s62, F32.mul(x7, w6)), F32.add(s63, F32.mul(x7, w7)), U32.is_lt(U32.add(j, 1), inner))
290
+
291
+ def packed_join(a: PackedBuffers, b: PackedBuffers) -> PackedBuffers:
292
+ JBuffers{xa, wa, oa} = a
293
+ JBuffers{xb, wb, ob} = b
294
+ JBuffers{Array.join(Input8, xa, xb), Array.join(Weight8, wa, wb), Array.join(Output64, oa, ob)}
295
+
296
+ @unsafe
297
+ def packed_step(x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>,
298
+ +lo: U32, +hi: U32, +tiles_col: U32, +inner: U32, active: Bool) -> PackedBuffers:
299
+ match active:
300
+ case False{}:
301
+ JBuffers{x, w, o}
302
+ case True{}:
303
+ row = U32.div(lo, tiles_col)
304
+ col = U32.mod(lo, tiles_col)
305
+ packed_dot(x, w, o, row, col, tiles_col, inner, hi, 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, True{})
306
+
307
+ @unsafe
308
+ def packed_split(x: Array<Input8> & Array<Input8>, w: Array<Weight8> & Array<Weight8>, o: Array<Output64> & Array<Output64>,
309
+ +lo: U32, +hi: U32, +tiles_col: U32, +grain: U32, +inner: U32) -> PackedBuffers:
310
+ (xa, xb) = x
311
+ (wa, wb) = w
312
+ (oa, ob) = o
313
+ +mid = U32.add(lo, U32.div(U32.sub(hi, lo), 2))
314
+ left right = packed_gemm(xa, wa, oa, lo, mid, tiles_col, grain, inner, U32.is_le(U32.sub(mid, lo), grain)) packed_gemm(xb, wb, ob, mid, hi, tiles_col, grain, inner, U32.is_le(U32.sub(hi, mid), grain))
315
+ packed_join(left, right)
316
+
317
+ @unsafe
318
+ def packed_gemm(x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>,
319
+ +lo: U32, +hi: U32, +tiles_col: U32, +grain: U32, +inner: U32, leaf: Bool) -> PackedBuffers:
320
+ match leaf:
321
+ case True{}:
322
+ row = U32.div(lo, tiles_col)
323
+ col = U32.mod(lo, tiles_col)
324
+ packed_dot(x, w, o, row, col, tiles_col, inner, hi, 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, True{})
325
+ case False{}:
326
+ packed_split(Array.fork(Input8, x), Array.fork(Weight8, w), Array.fork(Output64, o), lo, hi, tiles_col, grain, inner)
327
+
328
+ def packed_input() -> IO(Array<Input8>):
329
+ import "./bridge.c"
330
+
331
+ def packed_weights() -> IO(Array<Weight8>):
332
+ import "./bridge.c"
333
+
334
+ def packed_out() -> IO(Array<Output64>):
335
+ import "./bridge.c"
336
+
337
+ def packed_output(result: PackedBuffers) -> IO(Unit):
338
+ import "./bridge.c"
339
+
340
+ def main() -> IO(Unit):
341
+ do IO<Unit>:
342
+ +tree : Candidates <- input()
343
+ output(winner(tree))
344
+ output(softmax(tree))
345
+ vector : Features <- features()
346
+ write_features(layernorm(vector, 384.0, 0.00001))
347
+ rows : NormRows <- batch_features()
348
+ write_batch(norm_rows(rows, 384.0, 0.00001))
349
+ +tile : TileInput <- tile_input()
350
+ +w : JuliaMatrix <- dense_weights()
351
+ tile_output(dense_tile(tile, w))
352
+ tile_output(dense_tile(tile, w))
353
+ +tiles : InputTiles <- input_tiles()
354
+ +matrix : JuliaMatrix <- dense_weights()
355
+ output_tiles(dense_batch(tiles, matrix))
356
+ output_tiles(dense_batch(tiles, matrix))
357
+ px : Array<Input8> <- packed_input()
358
+ pw : Array<Weight8> <- packed_weights()
359
+ po : Array<Output64> <- packed_out()
360
+ packed_output(packed_gemm(px, pw, po, 0, 1, 1, 1, 1, True{}))
julia/router/router.py ADDED
@@ -0,0 +1,173 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Bounded Julia inference with a Bend candidate-tree reducer."""
2
+ from collections import OrderedDict
3
+ from dataclasses import dataclass
4
+ import json
5
+ import math
6
+ import threading
7
+
8
+ from ..probabilities import display_probabilities
9
+ from .native import BendReducer
10
+
11
+
12
+ @dataclass(frozen=True)
13
+ class RouteResult:
14
+ index: int
15
+ candidates: tuple[int, ...]
16
+ probabilities: tuple[float, ...]
17
+ rounds: int
18
+ model_rows: int
19
+ cache_hits: int
20
+ hierarchical: bool
21
+ # Probabilities are conditional on candidates, never global for a tournament.
22
+ probability_scope: str
23
+
24
+
25
+ class Router:
26
+ """Wrap an Engine-compatible logits(rows) implementation.
27
+
28
+ Up to width options: one unchanged model request. Larger choice requests:
29
+ retain survivors per group and rerank until one final group remains.
30
+ This increases supported option count, not the model's trained capacity.
31
+ """
32
+ def __init__(self, engine, *, library=None, width=20, survivors=2,
33
+ batch_size=16, cache_size=0, max_options=4096):
34
+ for name, value in [('width', width), ('survivors', survivors),
35
+ ('batch_size', batch_size), ('cache_size', cache_size),
36
+ ('max_options', max_options)]:
37
+ if type(value) is not int:
38
+ raise ValueError(f'{name} must be an integer')
39
+ if not 2 <= width <= 20 or not 1 <= survivors < width:
40
+ raise ValueError('Require 2 <= width <= 20 and 1 <= survivors < width')
41
+ if batch_size < 1 or cache_size < 0 or max_options < width:
42
+ raise ValueError('Invalid batch/cache/capacity limit')
43
+ self.engine, self.reducer = engine, BendReducer(library)
44
+ self.width, self.survivors = width, survivors
45
+ self.batch_size, self.cache_size, self.max_options = batch_size, cache_size, max_options
46
+ self._cache = OrderedDict()
47
+ self._lock = threading.RLock()
48
+
49
+ def clear_cache(self):
50
+ """Call after modifying model weights or inference settings."""
51
+ with self._lock:
52
+ self._cache.clear()
53
+
54
+ def _validate(self, row):
55
+ if not isinstance(row, dict):
56
+ raise ValueError('A request must be a dictionary')
57
+ if not isinstance(row.get('state'), (str, dict, list)) or not isinstance(row.get('question'), str):
58
+ raise ValueError('state must be text/JSON and question must be text')
59
+ options = row.get('options')
60
+ if not isinstance(options, list) or not 2 <= len(options) <= self.max_options:
61
+ raise ValueError(f'Expected 2–{self.max_options} options')
62
+ if not all(isinstance(x, str) and x for x in options):
63
+ raise ValueError('Options must be nonempty strings')
64
+ kind = row.get('type', 'choice')
65
+ if kind not in ('choice', 'score', 'noul'):
66
+ raise ValueError('Unknown decision type')
67
+ if kind == 'noul' and len(options) != 2:
68
+ raise ValueError('noul requires [false, true]')
69
+ if kind != 'choice' and len(options) > self.width:
70
+ raise ValueError('Hierarchical routing supports choice decisions only')
71
+ clean = dict(state=row['state'], question=row['question'], options=options, type=kind)
72
+ # Snapshot mutable state; preserve dict order used by Julia serialization.
73
+ return json.loads(json.dumps(clean, ensure_ascii=False, allow_nan=False))
74
+
75
+ def _score(self, jobs):
76
+ values = [None] * len(jobs)
77
+ missing = OrderedDict()
78
+ hits = 0
79
+ for i, row in enumerate(jobs):
80
+ key = json.dumps(row, ensure_ascii=False, separators=(',', ':'), allow_nan=False)
81
+ if key in self._cache:
82
+ values[i] = self._cache[key]
83
+ self._cache.move_to_end(key)
84
+ hits += 1
85
+ elif key in missing:
86
+ missing[key][1].append(i)
87
+ hits += 1
88
+ else:
89
+ missing[key] = (row, [i])
90
+ pending = list(missing.items())
91
+ for offset in range(0, len(pending), self.batch_size):
92
+ chunk = pending[offset:offset + self.batch_size]
93
+ output = list(self.engine.logits([entry[1][0] for entry in chunk]))
94
+ if len(output) != len(chunk):
95
+ raise ValueError('Engine returned the wrong number of rows')
96
+ for (key, (row, indices)), scores in zip(chunk, output):
97
+ scores = tuple(float(x) for x in scores)
98
+ if len(scores) != len(row['options']) or not all(math.isfinite(x) for x in scores):
99
+ raise ValueError('Engine logits must be finite and match option count')
100
+ for i in indices:
101
+ values[i] = scores
102
+ if self.cache_size:
103
+ self._cache[key] = scores
104
+ self._cache.move_to_end(key)
105
+ while len(self._cache) > self.cache_size:
106
+ self._cache.popitem(last=False)
107
+ return values, len(pending), hits
108
+
109
+ @staticmethod
110
+ def _confident_winner(scores, best):
111
+ """Keep one candidate when its local softmax dominates the group."""
112
+ runner_up = max(score for i, score in enumerate(scores) if i != best)
113
+ if runner_up - scores[best] >= math.log(0.045 / 0.95):
114
+ return False
115
+ total = sum(math.exp(score - scores[best]) for score in scores)
116
+ return 1 / total > 0.95 and math.exp(runner_up - scores[best]) / total < 0.045
117
+
118
+ def route(self, row):
119
+ return self.route_many([row])[0]
120
+
121
+ def route_many(self, rows):
122
+ """Batch independent groups across requests; preserve request/option order.
123
+
124
+ Calls on one Router serialize because most model engines and the LRU are
125
+ mutable. The C bridge also serializes access to Bend's global runtime.
126
+ """
127
+ with self._lock:
128
+ requests = [self._validate(row) for row in rows]
129
+ candidates = [list(range(len(row['options']))) for row in requests]
130
+ results = [None] * len(requests)
131
+ rounds = [0] * len(requests)
132
+ model_rows = hits = 0
133
+ while any(result is None for result in results):
134
+ jobs, layout = [], []
135
+ for i, row in enumerate(requests):
136
+ if results[i] is not None:
137
+ continue
138
+ rounds[i] += 1
139
+ current = candidates[i]
140
+ final = len(current) <= self.width
141
+ groups = [current[j:j + self.width] for j in range(0, len(current), self.width)]
142
+ candidates[i] = []
143
+ for group in groups:
144
+ if len(group) == 1:
145
+ candidates[i].extend(group)
146
+ continue
147
+ jobs.append(dict(row, options=[row['options'][k] for k in group]))
148
+ layout.append((i, group, final))
149
+ scored, used, cached = self._score(jobs)
150
+ model_rows += used
151
+ hits += cached
152
+ for (i, group, final), scores in zip(layout, scored):
153
+ best = self.reducer.argmax(scores)
154
+ if final:
155
+ maximum = max(scores)
156
+ weights = [math.exp(x - maximum) for x in scores]
157
+ total = sum(weights)
158
+ results[i] = RouteResult(group[best], tuple(group),
159
+ tuple(display_probabilities([x / total for x in weights])), rounds[i], 0, 0,
160
+ len(requests[i]['options']) > self.width,
161
+ 'final_candidates' if len(requests[i]['options']) > self.width else 'all_options')
162
+ elif self._confident_winner(scores, best):
163
+ candidates[i].append(group[best])
164
+ else:
165
+ remaining = list(range(len(group)))
166
+ chosen = []
167
+ for _ in range(min(self.survivors, len(group))):
168
+ position = self.reducer.argmax([scores[k] for k in remaining])
169
+ chosen.append(group[remaining.pop(position)])
170
+ candidates[i].extend(sorted(chosen))
171
+ # These counters are shared batch totals, not per-request attribution.
172
+ from dataclasses import replace
173
+ return [replace(result, model_rows=model_rows, cache_hits=hits) for result in results]
julia/router/specialize.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Bounds/alias specialization for the validated packed Bend CPU entry point.
2
+
3
+ Arithmetic and iteration order remain generated from router.bend. This pass is
4
+ valid only for bridge.c's distinct packed buffers and checked tile dimensions.
5
+ """
6
+ import re
7
+
8
+
9
+ def packed_cpu(source):
10
+ marker = ' WL_CASE(FID_PACKED_DOT)'
11
+ start = source.index(marker)
12
+ stop = source.index('#endif', start)
13
+ body = source[start:stop]
14
+ # Do not generalize this to Array.get: ordinary Bend arrays wrap indices.
15
+ # The bridge pads the full tile ranges and enforces positive bounded sizes.
16
+ specs = {'x': ('const ', 3), 'w': ('const ', 3), 'o': ('', 6)}
17
+ declarations = []
18
+ for name, (qualifier, shift) in specs.items():
19
+ var = f'_{name}_0'
20
+ locs = re.findall(rf'Term (_at_\d+) = blk_loc\(e.mem, {var}\);', body)
21
+ if len(locs) != 1:
22
+ raise RuntimeError(f'Unsupported packed {name} buffer lowering')
23
+ loc = locs[0]
24
+ declarations.append(f' {qualifier}u32a *restrict julia_{name} = '
25
+ f'({qualifier}u32a *)(e.mem + blk_loc(e.mem, {var}));')
26
+ body = body.replace(f'Term {loc} = blk_loc(e.mem, {var});', '')
27
+ pattern = rf'blk_at\({var}, (.+), {shift}\)'
28
+ body, count = re.subn(pattern, rf'((u32)(\1) << {shift})', body)
29
+ if count != 1:
30
+ raise RuntimeError(f'Unsupported packed {name} index lowering')
31
+ body = re.sub(rf'blk_read\(e.mem, 0, {loc}, ([^)]+)\)',
32
+ rf'julia_{name}[\1]', body)
33
+ body = re.sub(rf'blk_write\(e.mem, 0, {loc}, ([^,]+), ([^)]+)\);',
34
+ rf'julia_{name}[\1] = \2;', body)
35
+ if 'blk_at(' in body or 'blk_read(' in body or 'blk_write(' in body:
36
+ raise RuntimeError('Unrecognized packed buffer access; refusing specialization')
37
+ body = body.replace(' WL_OPEN', '\n'.join(declarations) + '\n WL_OPEN', 1)
38
+ return source[:start] + body + source[stop:]
julia/router/tests/synthetic.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Deterministic RANDOM-weight fixtures, never trained-quality evidence."""
2
+ from pathlib import Path
3
+
4
+
5
+ def make_checkpoint(directory, *, tiny=False):
6
+ import torch
7
+ from tokenizers import Tokenizer
8
+ from tokenizers.models import WordLevel
9
+ from tokenizers.pre_tokenizers import Whitespace
10
+ from transformers import PreTrainedTokenizerFast, ModernBertConfig, AutoModel
11
+ from julia.model import JuliaDecisionModel
12
+ root = Path(directory)
13
+ torch.manual_seed(42)
14
+ width, layers, heads = (64, 2, 2) if tiny else (384, 22, 6)
15
+ vocabulary = {'[PAD]': 0, '[CLS]': 1, '[SEP]': 2, '[UNK]': 3, '[MASK]': 4}
16
+ vocabulary.update({f'word{i}': i + 5 for i in range(507)})
17
+ backend = Tokenizer(WordLevel(vocabulary, unk_token='[UNK]'))
18
+ backend.pre_tokenizer = Whitespace()
19
+ tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, pad_token='[PAD]',
20
+ cls_token='[CLS]', sep_token='[SEP]', unk_token='[UNK]', mask_token='[MASK]')
21
+ tokenizer.save_pretrained(root / 'tokenizer')
22
+ config = ModernBertConfig(vocab_size=512, hidden_size=width,
23
+ intermediate_size=width * 3, num_hidden_layers=layers, num_attention_heads=heads,
24
+ max_position_embeddings=1024, local_attention=128, global_attn_every_n_layers=3,
25
+ pad_token_id=0, cls_token_id=1, sep_token_id=2, reference_compile=False)
26
+ encoder = AutoModel.from_config(config, attn_implementation='sdpa')
27
+ model = JuliaDecisionModel(encoder, head_layers=2)
28
+ model.save_pretrained(root)
29
+ return root
30
+
31
+
32
+ def sample_rows(count=16):
33
+ lengths = [8, 160, 24, 96]
34
+ return [dict(state=' '.join(f'word{(j + i) % 500}' for j in range(lengths[i % 4])),
35
+ question='Choose the matching option', type='choice',
36
+ options=[f'word{k + i}' for k in range(2 + i % 7)]) for i in range(count)]
julia/router/tests/test_engine.py ADDED
@@ -0,0 +1,367 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import tempfile
2
+ import unittest
3
+ from unittest.mock import patch
4
+ import numpy as np
5
+ import torch
6
+
7
+ from julia.data import Collator, sequence
8
+ from julia.inference import TransformerEngine as Engine
9
+ from julia.router.engine import FastEngine
10
+ from julia.router.native import BendReducer
11
+ from synthetic import make_checkpoint, sample_rows
12
+
13
+
14
+ class NativeMathTests(unittest.TestCase):
15
+ def test_transformer_math(self):
16
+ reducer = BendReducer()
17
+ rng = np.random.default_rng(42)
18
+ for width in [1, 7, 64, 384]:
19
+ x = rng.normal(size=(3, width)).astype('float32')
20
+ gamma = rng.normal(size=width).astype('float32')
21
+ beta = rng.normal(size=width).astype('float32')
22
+ expected = torch.nn.functional.layer_norm(torch.from_numpy(x), (width,),
23
+ torch.from_numpy(gamma), torch.from_numpy(beta)).numpy()
24
+ np.testing.assert_allclose(reducer.layernorm(x, gamma, beta), expected, atol=2e-6, rtol=2e-5)
25
+ for row in x:
26
+ index, probs = reducer.softmax(row)
27
+ self.assertEqual(index, int(row.argmax()))
28
+ np.testing.assert_allclose(probs, torch.from_numpy(row).softmax(-1).numpy(), atol=2e-7)
29
+ _, probabilities = reducer.softmax(np.array([10000, 10001, -10000], np.float32))
30
+ self.assertAlmostEqual(float(probabilities.sum()), 1.0)
31
+
32
+ def test_resident_dense_tiles_reuse_and_partial_edges(self):
33
+ from julia.router.native import BendMatrix
34
+ rng = np.random.default_rng(52)
35
+ for inputs, outputs, rows in [(1, 1, 1), (7, 5, 9), (64, 17, 7), (33, 65, 19), (384, 1152, 5)]:
36
+ weights = rng.normal(size=(outputs, inputs)).astype('float32')
37
+ values = rng.normal(size=(rows, inputs)).astype('float32')
38
+ matrix = BendMatrix(weights)
39
+ reference = values @ weights.T
40
+ first = matrix(values)
41
+ np.testing.assert_allclose(first, reference, atol=1e-4, rtol=1e-4)
42
+ for _ in range(8):
43
+ np.testing.assert_array_equal(matrix(values), first)
44
+ with self.assertRaises(ValueError): matrix(np.ones((2, inputs + 1), np.float32))
45
+ with self.assertRaises(ValueError): matrix(np.full((1, inputs), np.nan, np.float32))
46
+ matrix.close()
47
+ matrix.close()
48
+ with self.assertRaises(RuntimeError): matrix(values)
49
+ with self.assertRaises(ValueError): BendMatrix([[float('nan')]])
50
+
51
+ def test_tensor_chunks_and_opt_in_native_profile(self):
52
+ from julia.router.native import BendMatrix
53
+ reducer = BendReducer()
54
+ weights = np.random.default_rng(8).normal(size=(257, 3)).astype('float32')
55
+ matrix = BendMatrix(weights, reducer)
56
+ # Noncontiguous input and > 1M output cells require two native chunks.
57
+ values = torch.arange(4090 * 6, dtype=torch.float32).reshape(4090, 6)[:, ::2] / 1000
58
+ try:
59
+ reducer.profile(enabled=True, reset=True)
60
+ actual = matrix.tensor(values)
61
+ measured = reducer.profile(enabled=False)
62
+ self.assertEqual(measured['calls'], 2)
63
+ self.assertTrue(all(measured[k] >= 0 for k in ('pack_ms', 'execute_ms', 'unpack_ms')))
64
+ torch.testing.assert_close(actual, values @ torch.from_numpy(weights).T, rtol=1e-5, atol=1e-5)
65
+ matrix.tensor(values[:1])
66
+ self.assertEqual(reducer.profile(), measured)
67
+ with self.assertRaises(ValueError):
68
+ matrix.tensor(torch.full((1, 3), float('nan')))
69
+ finally:
70
+ reducer.profile(enabled=False, reset=True)
71
+ matrix.close()
72
+
73
+ def test_layernorm_parallel_callers_and_large_rows(self):
74
+ import concurrent.futures
75
+ reducer = BendReducer()
76
+ rng = np.random.default_rng(19)
77
+ values = [rng.normal(size=(rows, 384)).astype('float32') for rows in (1, 7, 32, 128)]
78
+ references = [torch.nn.functional.layer_norm(torch.from_numpy(x), (384,)).numpy() for x in values]
79
+ with concurrent.futures.ThreadPoolExecutor(4) as pool:
80
+ actual = list(pool.map(reducer.layernorm, values * 8))
81
+ for i, result in enumerate(actual):
82
+ np.testing.assert_allclose(result, references[i % len(values)], atol=3e-6, rtol=2e-5)
83
+
84
+ def test_packed_matrix_concurrent_owners(self):
85
+ from concurrent.futures import ThreadPoolExecutor
86
+ from julia.router.native import BendMatrix
87
+ rng = np.random.default_rng(2026)
88
+ weights = [rng.normal(size=(n, 33)).astype('float32') for n in (5, 17, 128)]
89
+ matrices = [BendMatrix(w) for w in weights]
90
+ values = rng.normal(size=(19, 33)).astype('float32')
91
+ try:
92
+ with ThreadPoolExecutor(6) as pool:
93
+ results = list(pool.map(lambda i: matrices[i % 3](values), range(96)))
94
+ for i, actual in enumerate(results):
95
+ np.testing.assert_allclose(actual, values @ weights[i % 3].T, atol=2e-5, rtol=2e-5)
96
+ finally:
97
+ for matrix in matrices:
98
+ matrix.close()
99
+
100
+ def test_rejected_native_math(self):
101
+ reducer = BendReducer()
102
+ for x in [[], [float('nan')], [float('inf')]]:
103
+ with self.assertRaises(ValueError):
104
+ reducer.softmax(x)
105
+ for kwargs in [dict(epsilon=0), dict(gamma=[1]), dict(beta=[0]), dict(epsilon=float('nan'))]:
106
+ with self.assertRaises(ValueError):
107
+ reducer.layernorm([[1, 2, 3]], **kwargs)
108
+
109
+
110
+ class ProbabilityDisplayTests(unittest.TestCase):
111
+ def test_rounds_decisive_answer_to_certain(self):
112
+ from julia.probabilities import display_probabilities
113
+
114
+ self.assertEqual(display_probabilities([0.955, 0.044, 0.001]), [1.0, 0.0, 0.0])
115
+ self.assertNotEqual(display_probabilities([0.951, 0.049]), [1.0, 0.0])
116
+
117
+ def test_redistributes_sub_one_percent_values(self):
118
+ from julia.probabilities import display_probabilities
119
+
120
+ result = display_probabilities([0.6, 0.395, 0.005])
121
+ self.assertEqual(result[2], 0.0)
122
+ self.assertAlmostEqual(sum(result), 1.0)
123
+ self.assertAlmostEqual(result[0] / result[1], 0.6 / 0.395)
124
+
125
+
126
+ class InputValidationTests(unittest.TestCase):
127
+ def test_public_loader_rejects_invalid_backend(self):
128
+ from julia import load_model
129
+ from julia.inference import Engine
130
+
131
+ self.assertIs(load_model, Engine)
132
+ with self.assertRaisesRegex(ValueError, 'backend'):
133
+ load_model('unused', backend='unknown')
134
+
135
+ def test_rejects_malformed_requests(self):
136
+ from julia.data import validate_row
137
+
138
+ valid = {'state': 'context', 'question': 'Choose', 'options': ['A', 'B']}
139
+ for row in (None, [], {**valid, 'teacher_logits': 'bad'},
140
+ {**valid, 'teacher_logits': [0.0, float('nan')]},
141
+ {**valid, 'teacher_logits': [True, 0.0]}):
142
+ with self.subTest(row=row), self.assertRaises(ValueError):
143
+ validate_row(row, 1)
144
+
145
+
146
+ class InferenceTests(unittest.TestCase):
147
+ @classmethod
148
+ def setUpClass(cls):
149
+ torch.set_num_threads(2)
150
+ cls.temp = tempfile.TemporaryDirectory()
151
+ cls.path = make_checkpoint(cls.temp.name, tiny=True)
152
+ cls.baseline = Engine(cls.path, device='cpu')
153
+ with patch('julia.router.native.BendReducer', side_effect=AssertionError('Default CPU must not load Bend')):
154
+ cls.fast = FastEngine(cls.path, device='cpu', batch_size=2)
155
+ cls.rows = sample_rows(5)
156
+ cls.rows[0]['state'] = {'text': 'Olá 世界 [MASK]', 'items': [1, 2]}
157
+
158
+ @classmethod
159
+ def tearDownClass(cls):
160
+ del cls.baseline, cls.fast
161
+ cls.temp.cleanup()
162
+
163
+ def test_modernbert_decision_encoder_deduplicates_rope(self):
164
+ from transformers import ModernBertConfig, ModernBertModel
165
+ from julia.model import JuliaDecisionModel
166
+ from julia.router.encoder import specialize_decision_encoder
167
+ config = ModernBertConfig(vocab_size=64, pad_token_id=0, bos_token_id=1, eos_token_id=2, hidden_size=32, intermediate_size=64,
168
+ num_hidden_layers=4, num_attention_heads=4, max_position_embeddings=128,
169
+ local_attention=16, global_attn_every_n_layers=3, attention_dropout=0.,
170
+ embedding_dropout=0., mlp_dropout=0.)
171
+ config._attn_implementation = 'sdpa'
172
+ model = JuliaDecisionModel(ModernBertModel(config), head_layers=0).eval()
173
+ ids = torch.arange(24).reshape(2,12) % 64
174
+ mask = torch.ones_like(ids); mask[1,8:] = 0
175
+ calls=[]
176
+ hook=model.encoder.rotary_emb.register_forward_hook(lambda *args: calls.append(1))
177
+ with torch.inference_mode():
178
+ expected=model.encoder(input_ids=ids,attention_mask=mask).last_hidden_state
179
+ self.assertEqual(len(calls), 4)
180
+ calls.clear()
181
+ self.assertTrue(specialize_decision_encoder(model))
182
+ actual=model.encoder(input_ids=ids,attention_mask=mask).last_hidden_state
183
+ self.assertEqual(len(calls), 2)
184
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
185
+ hook.remove()
186
+
187
+ def test_default_backend_and_padding_budget(self):
188
+ self.assertEqual(self.fast.transformer_backend, 'torch')
189
+ self.assertIsNone(self.fast.bend)
190
+ self.assertFalse(hasattr(self.fast, 'bend_projection_count'))
191
+ rows = [dict(self.rows[0], state='short'), dict(self.rows[1], state='long context ' * 100),
192
+ dict(self.rows[2], state='another short state')]
193
+ encoded = self.fast._encode(rows)
194
+ groups = list(self.fast._batch_indices(encoded))
195
+ self.assertEqual(sorted(i for group in groups for i in group), list(range(3)))
196
+ for group in groups:
197
+ lengths = [len(encoded[i]['ids']) for i in group]
198
+ self.assertLessEqual(max(lengths) * len(group), self.fast.padding_ratio * sum(lengths))
199
+ actual = self.fast.predict(rows)
200
+ expected = self.baseline.predict(rows)
201
+ for a, b in zip(actual, expected):
202
+ np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)
203
+
204
+ def test_lfs_pointer_error(self):
205
+ from pathlib import Path
206
+ with tempfile.TemporaryDirectory() as directory:
207
+ (Path(directory) / 'model.safetensors').write_text('version https://git-lfs.github.com/spec/v1\n')
208
+ with self.assertRaisesRegex(ValueError, 'Git LFS pointers'):
209
+ FastEngine(directory)
210
+
211
+ def test_serialization_and_pack(self):
212
+ encoded = self.fast._encode(self.rows)
213
+ self.assertEqual(encoded, [sequence(self.baseline.tokenizer, row) for row in self.rows])
214
+ packed = self.fast._pack(encoded)
215
+ reference = Collator(self.baseline.tokenizer, 1024, 256)(self.rows)
216
+ for name in packed:
217
+ self.assertTrue(torch.equal(packed[name], reference[name]), name)
218
+ self.assertEqual(self.fast.predict([]), [])
219
+ self.assertEqual(self.fast.logits([]), [])
220
+
221
+ def test_inference_collation_omits_training_tensors(self):
222
+ rows = [dict(row, target=0, teacher_logits=[0.] * len(row['options']))
223
+ for row in self.rows]
224
+ collate = self.baseline.collate
225
+ training = collate(rows)
226
+ inference = collate(rows, include_targets=False)
227
+ self.assertIn('labels', training)
228
+ self.assertIn('teacher_logits', training)
229
+ self.assertEqual(set(training) - set(inference), {'labels', 'teacher_logits'})
230
+ for name, value in inference.items():
231
+ self.assertTrue(torch.equal(value, training[name]), name)
232
+ self.assertEqual(self.baseline.predict(rows), self.baseline.predict(self.rows))
233
+
234
+ def test_nonfinite_real_scores_are_rejected(self):
235
+ original = self.fast.forward
236
+ try:
237
+ for invalid in (float('nan'), float('inf'), -float('inf')):
238
+ def forward(**batch):
239
+ result = torch.zeros_like(batch['marker_pos'], dtype=torch.float32)
240
+ result[0, 0] = invalid
241
+ return result
242
+ self.fast.forward = forward
243
+ with self.assertRaises(FloatingPointError):
244
+ self.fast.predict(self.rows)
245
+ finally:
246
+ self.fast.forward = original
247
+
248
+ def test_memory_mapped_checkpoint_dtype_and_logits(self):
249
+ from julia.model import JuliaDecisionModel
250
+ for dtype in (torch.float32, torch.bfloat16):
251
+ with tempfile.TemporaryDirectory() as directory:
252
+ model = JuliaDecisionModel.from_pretrained(self.path).to(dtype=dtype)
253
+ model.save_pretrained(directory)
254
+ loaded = JuliaDecisionModel.from_pretrained(directory, memory_map=True)
255
+ copied = JuliaDecisionModel.from_pretrained(directory, memory_map=False)
256
+ for name, value in loaded.state_dict().items():
257
+ self.assertEqual(value.dtype, copied.state_dict()[name].dtype)
258
+ torch.testing.assert_close(value, copied.state_dict()[name], rtol=0, atol=0)
259
+ if dtype == torch.float32:
260
+ batch = self.baseline.collate(self.rows)
261
+ with torch.inference_mode():
262
+ torch.testing.assert_close(loaded.eval()(**batch), copied.eval()(**batch), rtol=0, atol=0)
263
+
264
+ def test_selected_head_scores_actions_and_training(self):
265
+ model = self.baseline.model
266
+ batch = self.baseline.collate(self.rows)
267
+ batch['qtype'] = torch.arange(len(self.rows)) % 3
268
+ with torch.inference_mode():
269
+ model.marker_only_head = False
270
+ expected = model(**batch, return_actions=True)
271
+ model.marker_only_head = True
272
+ actual = model(**batch, return_actions=True)
273
+ for a, b in zip(expected, actual):
274
+ torch.testing.assert_close(a, b, atol=2e-6, rtol=2e-5)
275
+ # Training must retain the original full-sequence/dropout path.
276
+ model.train()
277
+ torch.manual_seed(123)
278
+ model.marker_only_head = False
279
+ expected = model(**batch)
280
+ torch.manual_seed(123)
281
+ model.marker_only_head = True
282
+ actual = model(**batch)
283
+ torch.testing.assert_close(expected, actual, atol=0, rtol=0)
284
+ actual.sum().backward()
285
+ self.assertTrue(torch.isfinite(model.scorer[-1].weight.grad).all())
286
+ model.zero_grad(set_to_none=True)
287
+ model.eval()
288
+ model.marker_only_head = False
289
+
290
+ def test_strict_encoding_reuses_audit_and_rejects_loss(self):
291
+ engine = FastEngine(self.path, device='cpu', strict_encoding=True)
292
+ row = dict(self.rows[1], state='word1 word2')
293
+ audit = engine.encoding_info([row])[0]
294
+ self.assertFalse(audit['stateTruncated'])
295
+ encoded = next(iter(engine._encoded.values()))
296
+ engine.predict([row])
297
+ self.assertIs(next(iter(engine._encoded.values())), encoded)
298
+ for bad in [dict(row, state='[MASK]'),
299
+ dict(row, state={'value': '[MASK]'}),
300
+ dict(row, options=['word1 ' * 49, 'word2']),
301
+ dict(row, state='word1 ' * 2000)]:
302
+ with self.assertRaises(ValueError):
303
+ engine.predict([bad])
304
+ # Changes in the encoding budget must not reuse stale cached IDs.
305
+ engine.max_length = 32
306
+ with self.assertRaises(ValueError):
307
+ engine.predict([dict(row, state='word1 ' * 100)])
308
+ engine.head_length = 16
309
+ with self.assertRaises(ValueError):
310
+ engine.predict([dict(row, question='word1 ' * 40)])
311
+
312
+ def test_inference_parity_and_order(self):
313
+ reference = self.baseline.predict(self.rows)
314
+ actual = self.fast.predict(self.rows)
315
+ self.assertEqual([x['index'] for x in reference], [x['index'] for x in actual])
316
+ for a, b in zip(reference, actual):
317
+ np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)
318
+ self.assertEqual(self.fast.predict(self.rows, probabilities=False),
319
+ [dict(index=x['index']) for x in actual])
320
+ self.assertTrue(self.fast._encoded)
321
+ self.fast.clear_cache()
322
+ self.assertFalse(self.fast._encoded)
323
+ self.assertFalse(self.fast._tokens.cache)
324
+
325
+ def test_bend_transformer_parity(self):
326
+ bend = FastEngine(self.path, device='cpu', batch_size=2,
327
+ transformer_backend='bend', bend_postprocess=True)
328
+ reference = self.baseline.predict(self.rows)
329
+ actual = bend.predict(self.rows)
330
+ # Different reduction orders can break an exact FP32 tie. Require the
331
+ # selected reference probability to be within tolerance of its maximum.
332
+ for a, b in zip(reference, actual):
333
+ p = a['probabilities']
334
+ self.assertLessEqual(max(p) - p[b['index']], 2e-6)
335
+ for a, b in zip(reference, actual):
336
+ np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)
337
+ self.assertEqual(bend.predict(self.rows, probabilities=False),
338
+ [dict(index=x['index']) for x in actual])
339
+
340
+ def test_bend_dense_encoder_parity(self):
341
+ dense = FastEngine(self.path, device='cpu', batch_size=2,
342
+ transformer_backend='bend-dense')
343
+ self.assertGreater(dense.bend_projection_count, 0)
344
+ expected = self.baseline.predict(self.rows)
345
+ actual = dense.predict(self.rows)
346
+ for a, b in zip(expected, actual):
347
+ np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)
348
+ self.assertLessEqual(max(a['probabilities']) - a['probabilities'][b['index']], 2e-6)
349
+ from julia.router.transformer import BendLinear
350
+ layer = next(m for m in dense.model.modules() if isinstance(m, BendLinear))
351
+ with torch.no_grad():
352
+ layer.weight.add_(1)
353
+ with self.assertRaisesRegex(RuntimeError, 'resident weights changed'):
354
+ dense.predict(self.rows)
355
+
356
+ @unittest.skipUnless(torch.cuda.is_available(), 'CUDA hardware unavailable')
357
+ def test_cuda(self):
358
+ baseline = Engine(self.path, device='cuda')
359
+ fast = FastEngine(self.path, device='cuda', batch_size=2)
360
+ reference, actual = baseline.predict(self.rows), fast.predict(self.rows)
361
+ self.assertEqual([x['index'] for x in reference], [x['index'] for x in actual])
362
+ for a, b in zip(reference, actual):
363
+ np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=5e-3)
364
+
365
+
366
+ if __name__ == '__main__':
367
+ unittest.main()
julia/router/tests/test_router.py ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import concurrent.futures
2
+ import math
3
+ import os
4
+ import random
5
+ import unittest
6
+
7
+ from julia.router import BendReducer, Router
8
+
9
+ LIB = os.environ.get('JULIA_ROUTER_LIBRARY', str(__import__('pathlib').Path(__file__).resolve().parents[1] / 'build/libjulia_router.so'))
10
+
11
+
12
+ class Engine:
13
+ def __init__(self):
14
+ self.batches = []
15
+
16
+ def logits(self, rows):
17
+ self.batches.append(rows)
18
+ return [[float(x) for x in row['options']] for row in rows]
19
+
20
+
21
+ def request(n):
22
+ return dict(state='Olá, 世界', question='Choose the largest number',
23
+ options=[str(i) for i in range(n)])
24
+
25
+
26
+ class Tests(unittest.TestCase):
27
+ def test_native_reference_and_reuse(self):
28
+ reducer = BendReducer(LIB)
29
+ rng = random.Random(10)
30
+ for n in [1, 2, 3, 20, 77, 1000]:
31
+ for _ in range(20):
32
+ values = [rng.randrange(-100, 100) for _ in range(n)]
33
+ self.assertEqual(reducer.argmax(values), max(range(n), key=values.__getitem__))
34
+ self.assertEqual(reducer.argmax([3, 3, -1]), 0)
35
+
36
+ def test_native_threads(self):
37
+ reducer = BendReducer(LIB)
38
+ with concurrent.futures.ThreadPoolExecutor(8) as pool:
39
+ self.assertEqual(list(pool.map(lambda _: reducer.argmax([1, 9, 2]), range(1000))), [1]*1000)
40
+
41
+ def test_invalid(self):
42
+ for scores in [[], [math.nan], [math.inf], [1e100]]:
43
+ with self.assertRaises(ValueError):
44
+ BendReducer(LIB).argmax(scores)
45
+ r = Router(Engine(), library=LIB)
46
+ for row in [request(1), dict(request(30), type='score'), dict(request(3), type='noul')]:
47
+ with self.assertRaises(ValueError):
48
+ r.route(row)
49
+
50
+ def test_direct_and_hierarchical(self):
51
+ engine = Engine()
52
+ router = Router(engine, library=LIB, batch_size=3)
53
+ results = router.route_many([request(n) for n in [2, 20, 21, 77, 4096]])
54
+ self.assertEqual([r.index for r in results], [1, 19, 20, 76, 4095])
55
+ for result in results:
56
+ self.assertAlmostEqual(sum(result.probabilities), 1)
57
+ self.assertIn(result.index, result.candidates)
58
+ self.assertEqual(results[0].probability_scope, 'all_options')
59
+ self.assertEqual(results[-1].probability_scope, 'final_candidates')
60
+ self.assertTrue(all(len(batch) <= 3 for batch in engine.batches))
61
+ self.assertTrue(all(2 <= len(row['options']) <= 20 for batch in engine.batches for row in batch))
62
+ self.assertEqual(router.route_many([]), [])
63
+
64
+ def test_confident_group_shortcut(self):
65
+ self.assertTrue(Router._confident_winner([0.0, 4.0, -3.0], 1))
66
+ self.assertFalse(Router._confident_winner([0.0, 3.0, -3.0], 1))
67
+
68
+ class DecisiveEngine(Engine):
69
+ def logits(self, rows):
70
+ self.batches.append(rows)
71
+ return [[10.0 * float(x) for x in row['options']] for row in rows]
72
+
73
+ engine = DecisiveEngine()
74
+ router = Router(engine, library=LIB)
75
+ result = router.route(request(2001))
76
+ self.assertEqual(result.index, 2000)
77
+ self.assertEqual(result.rounds, 3)
78
+ self.assertEqual(result.model_rows, 106)
79
+
80
+ baseline = Router(DecisiveEngine(), library=LIB)
81
+ baseline._confident_winner = lambda scores, best: False
82
+ original = baseline.route(request(2001))
83
+ self.assertEqual(original.index, result.index)
84
+ self.assertLess(result.rounds, original.rounds)
85
+ self.assertLess(result.model_rows, original.model_rows)
86
+
87
+ def test_cache_and_dedup(self):
88
+ engine = Engine()
89
+ router = Router(engine, library=LIB, cache_size=1)
90
+ results = router.route_many([request(5), request(5)])
91
+ self.assertEqual(results[0].model_rows, 1)
92
+ self.assertEqual(results[0].cache_hits, 1)
93
+ self.assertEqual(router.route(request(5)).model_rows, 0)
94
+ router.clear_cache()
95
+ self.assertEqual(router.route(request(5)).model_rows, 1)
96
+ router.route(request(6))
97
+ self.assertEqual(router.route(request(5)).model_rows, 1)
98
+
99
+ def test_engine_failure(self):
100
+ class Bad:
101
+ def logits(self, rows):
102
+ return [[0]] * len(rows)
103
+ with self.assertRaises(ValueError):
104
+ Router(Bad(), library=LIB).route(request(4))
105
+
106
+
107
+ if __name__ == '__main__':
108
+ unittest.main()
julia/router/transformer.py ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Resident Bend encoder projections and optional decision-head normalization.
2
+
3
+ The packed FP32 kernels own matrix arithmetic and coarse parallel scheduling.
4
+ PyTorch supplies embedding lookup, SDPA, activations and the selected-output head.
5
+ """
6
+ import numpy as np
7
+ import torch
8
+ from torch import nn
9
+
10
+
11
+ class BendLayerNorm(nn.Module):
12
+ def __init__(self, original, reducer):
13
+ super().__init__()
14
+ self.weight, self.bias = original.weight, original.bias
15
+ self.eps = original.eps
16
+ self.normalized_shape = original.normalized_shape
17
+ self.reducer = reducer
18
+
19
+ def forward(self, x):
20
+ if x.device.type != 'cpu' or x.dtype != torch.float32:
21
+ raise ValueError('Bend LayerNorm requires CPU float32; use torch backend for CUDA')
22
+ if torch.is_grad_enabled():
23
+ raise RuntimeError('Bend transformer kernels require torch.inference_mode()')
24
+ width = x.shape[-1]
25
+ values = x.detach().contiguous().numpy().reshape(-1, width)
26
+ gamma = self.weight.detach().numpy() if self.weight is not None else None
27
+ beta = self.bias.detach().numpy() if self.bias is not None else None
28
+ parts = [self.reducer.layernorm(values[start:start + 1048576 // width], gamma, beta, self.eps)
29
+ for start in range(0, len(values), 1048576 // width)]
30
+ result = parts[0] if len(parts) == 1 else np.concatenate(parts)
31
+ return torch.from_numpy(result.reshape(x.shape))
32
+
33
+
34
+ class BendHeadLayer(nn.Module):
35
+ """Explicit pre-norm block; prevents fused MHA from bypassing Bend norms."""
36
+ def __init__(self, layer, reducer):
37
+ super().__init__()
38
+ if not layer.norm_first:
39
+ raise ValueError('Julia expects pre-norm transformer layers')
40
+ self.layer = layer
41
+ self.norm1 = BendLayerNorm(layer.norm1, reducer)
42
+ self.norm2 = BendLayerNorm(layer.norm2, reducer)
43
+
44
+ def forward(self, x, src_key_padding_mask=None):
45
+ normalized = self.norm1(x)
46
+ attention = self.layer.self_attn(normalized, normalized, normalized,
47
+ key_padding_mask=src_key_padding_mask, need_weights=False)[0]
48
+ x = x + self.layer.dropout1(attention)
49
+ hidden = self.layer.linear2(self.layer.dropout(
50
+ self.layer.activation(self.layer.linear1(self.norm2(x)))))
51
+ return x + self.layer.dropout2(hidden)
52
+
53
+
54
+ def install_bend_head(model, reducer):
55
+ if model.head is not None:
56
+ model.head.layers = nn.ModuleList([BendHeadLayer(layer, reducer) for layer in model.head.layers])
57
+ model.scorer[0] = BendLayerNorm(model.scorer[0], reducer)
58
+ model.eval()
59
+
60
+
61
+ class BendLinear(nn.Module):
62
+ """Inference-only CPU projection using resident weights and Bend arithmetic."""
63
+ def __init__(self, original, reducer):
64
+ super().__init__()
65
+ from .native import BendMatrix
66
+ self.weight, self.bias = original.weight, original.bias
67
+ self.in_features, self.out_features = original.in_features, original.out_features
68
+ self.matrix = BendMatrix(original.weight.detach().numpy(), reducer)
69
+ self._weight_version = self.weight._version
70
+ self._weight_pointer = self.weight.data_ptr()
71
+
72
+ def _check_weight(self):
73
+ if self.weight._version != self._weight_version or self.weight.data_ptr() != self._weight_pointer:
74
+ raise RuntimeError('Bend resident weights changed; recreate the inference backend')
75
+
76
+ def forward(self, x):
77
+ if x.device.type != 'cpu' or x.dtype != torch.float32 or torch.is_grad_enabled():
78
+ raise ValueError('Bend dense projections require CPU FP32 inference_mode')
79
+ self._check_weight()
80
+ result = self.matrix.tensor(x, validate=False)
81
+ return result if self.bias is None else result + self.bias
82
+
83
+
84
+ def install_bend_encoder(model, reducer):
85
+ """Route all encoder linear projections through Bend; attention remains SDPA."""
86
+ count = 0
87
+ def install(module):
88
+ nonlocal count
89
+ for name, child in list(module.named_children()):
90
+ if type(child) is nn.Linear:
91
+ setattr(module, name, BendLinear(child, reducer))
92
+ count += 1
93
+ else:
94
+ install(child)
95
+ install(model.encoder)
96
+ model.eval()
97
+ return count
julia_config.json ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "format_version": 1,
3
+ "architecture": "JuliaDecisionModel",
4
+ "weight_dtype": "float32",
5
+ "head_layers": 2,
6
+ "n_act": 2,
7
+ "dropout": 0.1
8
+ }
metrics/accuracy-20260924.json ADDED
@@ -0,0 +1,434 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "date": "2026-09-24",
3
+ "weights_sha256": "df853bf7fe424420011f3d0c47a05d7341aa9eefa7fb9f203ea4aada4ad95b72",
4
+ "benchmarks": {
5
+ "typed-decisions": {
6
+ "count": 2000,
7
+ "correct": 1463,
8
+ "answered": 2000,
9
+ "accuracy": 0.7315,
10
+ "coverage": 1.0,
11
+ "jev_reference": 0.727,
12
+ "delta": 0.0045000000000000595,
13
+ "statuses": {
14
+ "ok": 2000
15
+ }
16
+ },
17
+ "agnews": {
18
+ "count": 100,
19
+ "correct": 94,
20
+ "answered": 100,
21
+ "accuracy": 0.94,
22
+ "coverage": 1.0,
23
+ "jev_reference": 0.91,
24
+ "delta": 0.029999999999999916,
25
+ "statuses": {
26
+ "ok": 100
27
+ }
28
+ },
29
+ "emotiondair": {
30
+ "count": 100,
31
+ "correct": 86,
32
+ "answered": 100,
33
+ "accuracy": 0.86,
34
+ "coverage": 1.0,
35
+ "jev_reference": 0.48,
36
+ "delta": 0.38,
37
+ "statuses": {
38
+ "ok": 100
39
+ }
40
+ },
41
+ "banking77": {
42
+ "count": 100,
43
+ "correct": 64,
44
+ "answered": 99,
45
+ "accuracy": 0.64,
46
+ "coverage": 0.99,
47
+ "jev_reference": 0.87,
48
+ "delta": -0.22999999999999998,
49
+ "statuses": {
50
+ "ok": 99,
51
+ "abstained": 1
52
+ }
53
+ }
54
+ },
55
+ "typed": {
56
+ "choice": {
57
+ "count": 600,
58
+ "correct": 428,
59
+ "accuracy": 0.7133333333333334
60
+ },
61
+ "noul": {
62
+ "count": 600,
63
+ "correct": 484,
64
+ "accuracy": 0.8066666666666666
65
+ },
66
+ "score": {
67
+ "count": 800,
68
+ "correct": 551,
69
+ "accuracy": 0.68875
70
+ }
71
+ },
72
+ "massive": {
73
+ "per_language": {
74
+ "mn-MN": {
75
+ "count": 2974,
76
+ "correct": 1618,
77
+ "unsupported": 0,
78
+ "accuracy": 0.5440484196368527
79
+ },
80
+ "af-ZA": {
81
+ "count": 2974,
82
+ "correct": 2253,
83
+ "unsupported": 0,
84
+ "accuracy": 0.7575655682582381
85
+ },
86
+ "el-GR": {
87
+ "count": 2974,
88
+ "correct": 2244,
89
+ "unsupported": 0,
90
+ "accuracy": 0.7545393409549428
91
+ },
92
+ "ta-IN": {
93
+ "count": 2974,
94
+ "correct": 1840,
95
+ "unsupported": 0,
96
+ "accuracy": 0.6186953597848016
97
+ },
98
+ "ar-SA": {
99
+ "count": 2974,
100
+ "correct": 2118,
101
+ "unsupported": 0,
102
+ "accuracy": 0.7121721587088097
103
+ },
104
+ "ca-ES": {
105
+ "count": 2974,
106
+ "correct": 2259,
107
+ "unsupported": 0,
108
+ "accuracy": 0.7595830531271015
109
+ },
110
+ "ur-PK": {
111
+ "count": 2974,
112
+ "correct": 2057,
113
+ "unsupported": 0,
114
+ "accuracy": 0.691661062542031
115
+ },
116
+ "pl-PL": {
117
+ "count": 2974,
118
+ "correct": 2377,
119
+ "unsupported": 0,
120
+ "accuracy": 0.7992602555480834
121
+ },
122
+ "ko-KR": {
123
+ "count": 2974,
124
+ "correct": 2296,
125
+ "unsupported": 0,
126
+ "accuracy": 0.7720242098184263
127
+ },
128
+ "az-AZ": {
129
+ "count": 2974,
130
+ "correct": 2166,
131
+ "unsupported": 0,
132
+ "accuracy": 0.7283120376597175
133
+ },
134
+ "da-DK": {
135
+ "count": 2974,
136
+ "correct": 2416,
137
+ "unsupported": 0,
138
+ "accuracy": 0.812373907195696
139
+ },
140
+ "kn-IN": {
141
+ "count": 2974,
142
+ "correct": 1654,
143
+ "unsupported": 0,
144
+ "accuracy": 0.5561533288500337
145
+ },
146
+ "tl-PH": {
147
+ "count": 2974,
148
+ "correct": 2146,
149
+ "unsupported": 0,
150
+ "accuracy": 0.7215870880968392
151
+ },
152
+ "is-IS": {
153
+ "count": 2974,
154
+ "correct": 1961,
155
+ "unsupported": 0,
156
+ "accuracy": 0.6593813046402152
157
+ },
158
+ "lv-LV": {
159
+ "count": 2974,
160
+ "correct": 1960,
161
+ "unsupported": 0,
162
+ "accuracy": 0.6590450571620713
163
+ },
164
+ "it-IT": {
165
+ "count": 2974,
166
+ "correct": 2432,
167
+ "unsupported": 0,
168
+ "accuracy": 0.8177538668459986
169
+ },
170
+ "es-ES": {
171
+ "count": 2974,
172
+ "correct": 2531,
173
+ "unsupported": 0,
174
+ "accuracy": 0.8510423671822461
175
+ },
176
+ "fr-FR": {
177
+ "count": 2974,
178
+ "correct": 2515,
179
+ "unsupported": 0,
180
+ "accuracy": 0.8456624075319435
181
+ },
182
+ "ml-IN": {
183
+ "count": 2974,
184
+ "correct": 1653,
185
+ "unsupported": 0,
186
+ "accuracy": 0.5558170813718897
187
+ },
188
+ "km-KH": {
189
+ "count": 2974,
190
+ "correct": 1421,
191
+ "unsupported": 0,
192
+ "accuracy": 0.47780766644250167
193
+ },
194
+ "fa-IR": {
195
+ "count": 2974,
196
+ "correct": 2382,
197
+ "unsupported": 0,
198
+ "accuracy": 0.800941492938803
199
+ },
200
+ "sw-KE": {
201
+ "count": 2974,
202
+ "correct": 1655,
203
+ "unsupported": 0,
204
+ "accuracy": 0.5564895763281775
205
+ },
206
+ "en-US": {
207
+ "count": 2974,
208
+ "correct": 2580,
209
+ "unsupported": 0,
210
+ "accuracy": 0.867518493611298
211
+ },
212
+ "tr-TR": {
213
+ "count": 2974,
214
+ "correct": 2388,
215
+ "unsupported": 0,
216
+ "accuracy": 0.8029589778076665
217
+ },
218
+ "bn-BD": {
219
+ "count": 2974,
220
+ "correct": 1972,
221
+ "unsupported": 0,
222
+ "accuracy": 0.6630800268997983
223
+ },
224
+ "he-IL": {
225
+ "count": 2974,
226
+ "correct": 2046,
227
+ "unsupported": 0,
228
+ "accuracy": 0.6879623402824478
229
+ },
230
+ "te-IN": {
231
+ "count": 2974,
232
+ "correct": 1900,
233
+ "unsupported": 0,
234
+ "accuracy": 0.6388702084734365
235
+ },
236
+ "pt-PT": {
237
+ "count": 2974,
238
+ "correct": 2565,
239
+ "unsupported": 0,
240
+ "accuracy": 0.8624747814391392
241
+ },
242
+ "ka-GE": {
243
+ "count": 2974,
244
+ "correct": 1501,
245
+ "unsupported": 0,
246
+ "accuracy": 0.5047074646940148
247
+ },
248
+ "ja-JP": {
249
+ "count": 2974,
250
+ "correct": 2446,
251
+ "unsupported": 0,
252
+ "accuracy": 0.8224613315400134
253
+ },
254
+ "id-ID": {
255
+ "count": 2974,
256
+ "correct": 2415,
257
+ "unsupported": 0,
258
+ "accuracy": 0.8120376597175522
259
+ },
260
+ "ru-RU": {
261
+ "count": 2974,
262
+ "correct": 2450,
263
+ "unsupported": 0,
264
+ "accuracy": 0.8238063214525891
265
+ },
266
+ "hy-AM": {
267
+ "count": 2974,
268
+ "correct": 1818,
269
+ "unsupported": 0,
270
+ "accuracy": 0.6112979152656355
271
+ },
272
+ "nb-NO": {
273
+ "count": 2974,
274
+ "correct": 2391,
275
+ "unsupported": 0,
276
+ "accuracy": 0.8039677202420982
277
+ },
278
+ "ms-MY": {
279
+ "count": 2974,
280
+ "correct": 2306,
281
+ "unsupported": 0,
282
+ "accuracy": 0.7753866845998655
283
+ },
284
+ "sq-AL": {
285
+ "count": 2974,
286
+ "correct": 2035,
287
+ "unsupported": 0,
288
+ "accuracy": 0.6842636180228648
289
+ },
290
+ "sv-SE": {
291
+ "count": 2974,
292
+ "correct": 2362,
293
+ "unsupported": 0,
294
+ "accuracy": 0.7942165433759247
295
+ },
296
+ "fi-FI": {
297
+ "count": 2974,
298
+ "correct": 2096,
299
+ "unsupported": 0,
300
+ "accuracy": 0.7047747141896435
301
+ },
302
+ "th-TH": {
303
+ "count": 2974,
304
+ "correct": 2299,
305
+ "unsupported": 0,
306
+ "accuracy": 0.7730329522528581
307
+ },
308
+ "de-DE": {
309
+ "count": 2974,
310
+ "correct": 2426,
311
+ "unsupported": 0,
312
+ "accuracy": 0.8157363819771352
313
+ },
314
+ "vi-VN": {
315
+ "count": 2974,
316
+ "correct": 2224,
317
+ "unsupported": 0,
318
+ "accuracy": 0.7478143913920645
319
+ },
320
+ "my-MM": {
321
+ "count": 2974,
322
+ "correct": 1513,
323
+ "unsupported": 0,
324
+ "accuracy": 0.5087424344317417
325
+ },
326
+ "jv-ID": {
327
+ "count": 2974,
328
+ "correct": 1802,
329
+ "unsupported": 0,
330
+ "accuracy": 0.6059179556153329
331
+ },
332
+ "sl-SL": {
333
+ "count": 2974,
334
+ "correct": 2147,
335
+ "unsupported": 0,
336
+ "accuracy": 0.7219233355749832
337
+ },
338
+ "hi-IN": {
339
+ "count": 2974,
340
+ "correct": 2237,
341
+ "unsupported": 0,
342
+ "accuracy": 0.7521856086079355
343
+ },
344
+ "nl-NL": {
345
+ "count": 2974,
346
+ "correct": 2449,
347
+ "unsupported": 0,
348
+ "accuracy": 0.8234700739744452
349
+ },
350
+ "zh-CN": {
351
+ "count": 2974,
352
+ "correct": 2487,
353
+ "unsupported": 0,
354
+ "accuracy": 0.8362474781439139
355
+ },
356
+ "hu-HU": {
357
+ "count": 2974,
358
+ "correct": 2161,
359
+ "unsupported": 0,
360
+ "accuracy": 0.726630800268998
361
+ },
362
+ "ro-RO": {
363
+ "count": 2974,
364
+ "correct": 2265,
365
+ "unsupported": 0,
366
+ "accuracy": 0.761600537995965
367
+ },
368
+ "zh-TW": {
369
+ "count": 2974,
370
+ "correct": 2365,
371
+ "unsupported": 0,
372
+ "accuracy": 0.7952252858103565
373
+ },
374
+ "cy-GB": {
375
+ "count": 2974,
376
+ "correct": 1639,
377
+ "unsupported": 0,
378
+ "accuracy": 0.551109616677875
379
+ },
380
+ "am-ET": {
381
+ "count": 2974,
382
+ "correct": 1334,
383
+ "unsupported": 0,
384
+ "accuracy": 0.4485541358439812
385
+ }
386
+ },
387
+ "macro_accuracy": 0.7149979307847499,
388
+ "count": 154648,
389
+ "seconds": 192.6575419740002,
390
+ "examples_per_second": 802.7092965863246
391
+ },
392
+ "diagnostics": {
393
+ "diagnostic:pt": {
394
+ "count": 3,
395
+ "accuracy": 1.0,
396
+ "worst_order_accuracy": 1.0,
397
+ "consistency": 1.0,
398
+ "brier_regret": 7.886988304767776e-10,
399
+ "worst_order_brier": 2.467773398956191e-09,
400
+ "nll": 2.103891552285609e-05,
401
+ "ece": 2.1038666961015196e-05
402
+ },
403
+ "diagnostic:en": {
404
+ "count": 3,
405
+ "accuracy": 1.0,
406
+ "worst_order_accuracy": 1.0,
407
+ "consistency": 1.0,
408
+ "brier_regret": 1.1595274914276354e-11,
409
+ "worst_order_brier": 5.446226103608541e-11,
410
+ "nll": 2.528017219883227e-06,
411
+ "ece": 2.528013180791433e-06
412
+ },
413
+ "diagnostic-alias:pt": {
414
+ "count": 2,
415
+ "accuracy": 1.0,
416
+ "worst_order_accuracy": 1.0,
417
+ "consistency": 1.0,
418
+ "brier_regret": 2.7625965108563243e-06,
419
+ "worst_order_brier": 5.561015463157067e-06,
420
+ "nll": 0.000956396971946862,
421
+ "ece": 0.0009555747324743735
422
+ },
423
+ "diagnostic-alias:en": {
424
+ "count": 2,
425
+ "accuracy": 1.0,
426
+ "worst_order_accuracy": 1.0,
427
+ "consistency": 1.0,
428
+ "brier_regret": 1.0381946113714448e-06,
429
+ "worst_order_brier": 2.0757069438387377e-06,
430
+ "nll": 0.0005331441533180119,
431
+ "ece": 0.0005328709204919235
432
+ }
433
+ }
434
+ }
metrics/validation.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "count": 3814,
3
+ "accuracy": 0.6982170949134766,
4
+ "nll": 0.7055923721868408,
5
+ "per_source": {
6
+ "agnews": {
7
+ "count": 505,
8
+ "accuracy": 0.8792079207920792,
9
+ "nll": 0.34157273733644883
10
+ },
11
+ "banking77": {
12
+ "count": 476,
13
+ "accuracy": 0.9495798319327731,
14
+ "nll": 0.13563624952062953
15
+ },
16
+ "emotion": {
17
+ "count": 483,
18
+ "accuracy": 0.9130434782608695,
19
+ "nll": 0.27308556848896093
20
+ },
21
+ "gsm8k": {
22
+ "count": 404,
23
+ "accuracy": 1.0,
24
+ "nll": 0.0004848763044099426
25
+ },
26
+ "hellaswag": {
27
+ "count": 489,
28
+ "accuracy": 0.36400817995910023,
29
+ "nll": 1.406457885786678
30
+ },
31
+ "mmlu": {
32
+ "count": 958,
33
+ "accuracy": 0.48643006263048016,
34
+ "nll": 1.342455263134215
35
+ },
36
+ "winogrande": {
37
+ "count": 499,
38
+ "accuracy": 0.5571142284569138,
39
+ "nll": 0.6976872796682789
40
+ }
41
+ },
42
+ "step": 900,
43
+ "checkpoint": "runs/julia-auto-1000/checkpoint-00000900",
44
+ "validation_sha256": "8f089f6bed53b85e9a839ce0210a867ebfa108f63bb2200bd291400d0af9d9f5"
45
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:df853bf7fe424420011f3d0c47a05d7341aa9eefa7fb9f203ea4aada4ad95b72
3
+ size 577189056
provenance.json ADDED
@@ -0,0 +1,527 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "family": "Julia-1",
3
+ "variant": "posttrained-candidate",
4
+ "parent_weights_sha256": "e0fce3f70bc1add0fc5ea1555728129775f313dd8beb40bb4d0284b9cc4582e0",
5
+ "weights_sha256": "df853bf7fe424420011f3d0c47a05d7341aa9eefa7fb9f203ea4aada4ad95b72",
6
+ "step": 500,
7
+ "validation": {
8
+ "agnews": {
9
+ "count": 500,
10
+ "correct": 461,
11
+ "accuracy": 0.922,
12
+ "nll": 0.2361866206965157
13
+ },
14
+ "banking-joint": {
15
+ "count": 500,
16
+ "correct": 453,
17
+ "accuracy": 0.906,
18
+ "nll": 0.4732417627487641
19
+ },
20
+ "banking-rank": {
21
+ "count": 500,
22
+ "correct": 483,
23
+ "accuracy": 0.966,
24
+ "nll": 0.12442472780195588
25
+ },
26
+ "bilingual:en": {
27
+ "count": 409,
28
+ "correct": 396,
29
+ "accuracy": 0.9682151589242054,
30
+ "nll": 0.1019394528733742
31
+ },
32
+ "bilingual:pt": {
33
+ "count": 409,
34
+ "correct": 399,
35
+ "accuracy": 0.9755501222493888,
36
+ "nll": 0.07665246521243735
37
+ },
38
+ "emotiondair": {
39
+ "count": 500,
40
+ "correct": 457,
41
+ "accuracy": 0.914,
42
+ "nll": 0.4869592022011608
43
+ },
44
+ "jev:arc_challenge": {
45
+ "count": 200,
46
+ "correct": 57,
47
+ "accuracy": 0.285,
48
+ "nll": 4.898485339760228
49
+ },
50
+ "jev:mmlu": {
51
+ "count": 198,
52
+ "correct": 52,
53
+ "accuracy": 0.26262626262626265,
54
+ "nll": 5.022009487458182
55
+ },
56
+ "jev:mnli": {
57
+ "count": 200,
58
+ "correct": 121,
59
+ "accuracy": 0.605,
60
+ "nll": 2.405353136634182
61
+ },
62
+ "jev:sst5": {
63
+ "count": 200,
64
+ "correct": 81,
65
+ "accuracy": 0.405,
66
+ "nll": 3.058030146173687
67
+ },
68
+ "massive:af-ZA": {
69
+ "count": 200,
70
+ "correct": 160,
71
+ "accuracy": 0.8,
72
+ "nll": 0.7254367060870702
73
+ },
74
+ "massive:am-ET": {
75
+ "count": 200,
76
+ "correct": 91,
77
+ "accuracy": 0.455,
78
+ "nll": 1.8078553098358898
79
+ },
80
+ "massive:ar-SA": {
81
+ "count": 200,
82
+ "correct": 143,
83
+ "accuracy": 0.715,
84
+ "nll": 0.9127131833431383
85
+ },
86
+ "massive:az-AZ": {
87
+ "count": 200,
88
+ "correct": 150,
89
+ "accuracy": 0.75,
90
+ "nll": 0.8805768398220111
91
+ },
92
+ "massive:bn-BD": {
93
+ "count": 200,
94
+ "correct": 135,
95
+ "accuracy": 0.675,
96
+ "nll": 1.0898693716287062
97
+ },
98
+ "massive:ca-ES": {
99
+ "count": 200,
100
+ "correct": 149,
101
+ "accuracy": 0.745,
102
+ "nll": 0.9042317060852588
103
+ },
104
+ "massive:cy-GB": {
105
+ "count": 200,
106
+ "correct": 107,
107
+ "accuracy": 0.535,
108
+ "nll": 1.5896150314416113
109
+ },
110
+ "massive:da-DK": {
111
+ "count": 200,
112
+ "correct": 155,
113
+ "accuracy": 0.775,
114
+ "nll": 0.7864874714409028
115
+ },
116
+ "massive:de-DE": {
117
+ "count": 200,
118
+ "correct": 163,
119
+ "accuracy": 0.815,
120
+ "nll": 0.6609104527319621
121
+ },
122
+ "massive:el-GR": {
123
+ "count": 200,
124
+ "correct": 146,
125
+ "accuracy": 0.73,
126
+ "nll": 0.9561585124287475
127
+ },
128
+ "massive:en-US": {
129
+ "count": 200,
130
+ "correct": 171,
131
+ "accuracy": 0.855,
132
+ "nll": 0.4239840084077537
133
+ },
134
+ "massive:es-ES": {
135
+ "count": 200,
136
+ "correct": 170,
137
+ "accuracy": 0.85,
138
+ "nll": 0.4821026878847341
139
+ },
140
+ "massive:fa-IR": {
141
+ "count": 200,
142
+ "correct": 153,
143
+ "accuracy": 0.765,
144
+ "nll": 0.7513568875009335
145
+ },
146
+ "massive:fi-FI": {
147
+ "count": 200,
148
+ "correct": 136,
149
+ "accuracy": 0.68,
150
+ "nll": 1.0470849014139052
151
+ },
152
+ "massive:fr-FR": {
153
+ "count": 200,
154
+ "correct": 172,
155
+ "accuracy": 0.86,
156
+ "nll": 0.5079110052088927
157
+ },
158
+ "massive:he-IL": {
159
+ "count": 200,
160
+ "correct": 137,
161
+ "accuracy": 0.685,
162
+ "nll": 1.1409396331540986
163
+ },
164
+ "massive:hi-IN": {
165
+ "count": 200,
166
+ "correct": 148,
167
+ "accuracy": 0.74,
168
+ "nll": 0.8160677297670518
169
+ },
170
+ "massive:hu-HU": {
171
+ "count": 200,
172
+ "correct": 142,
173
+ "accuracy": 0.71,
174
+ "nll": 0.9946744119686869
175
+ },
176
+ "massive:hy-AM": {
177
+ "count": 200,
178
+ "correct": 117,
179
+ "accuracy": 0.585,
180
+ "nll": 1.3989824168977554
181
+ },
182
+ "massive:id-ID": {
183
+ "count": 200,
184
+ "correct": 162,
185
+ "accuracy": 0.81,
186
+ "nll": 0.7148022617771757
187
+ },
188
+ "massive:is-IS": {
189
+ "count": 200,
190
+ "correct": 140,
191
+ "accuracy": 0.7,
192
+ "nll": 1.0531424324651082
193
+ },
194
+ "massive:it-IT": {
195
+ "count": 200,
196
+ "correct": 152,
197
+ "accuracy": 0.76,
198
+ "nll": 0.7013908650180337
199
+ },
200
+ "massive:ja-JP": {
201
+ "count": 200,
202
+ "correct": 157,
203
+ "accuracy": 0.785,
204
+ "nll": 0.7928279227610233
205
+ },
206
+ "massive:jv-ID": {
207
+ "count": 200,
208
+ "correct": 124,
209
+ "accuracy": 0.62,
210
+ "nll": 1.349707674896094
211
+ },
212
+ "massive:ka-GE": {
213
+ "count": 200,
214
+ "correct": 96,
215
+ "accuracy": 0.48,
216
+ "nll": 1.7132205322599072
217
+ },
218
+ "massive:km-KH": {
219
+ "count": 200,
220
+ "correct": 96,
221
+ "accuracy": 0.48,
222
+ "nll": 1.6912345944473137
223
+ },
224
+ "massive:kn-IN": {
225
+ "count": 200,
226
+ "correct": 122,
227
+ "accuracy": 0.61,
228
+ "nll": 1.5155959311408833
229
+ },
230
+ "massive:ko-KR": {
231
+ "count": 200,
232
+ "correct": 158,
233
+ "accuracy": 0.79,
234
+ "nll": 0.677267845838551
235
+ },
236
+ "massive:lv-LV": {
237
+ "count": 200,
238
+ "correct": 130,
239
+ "accuracy": 0.65,
240
+ "nll": 1.0922244557466885
241
+ },
242
+ "massive:ml-IN": {
243
+ "count": 200,
244
+ "correct": 113,
245
+ "accuracy": 0.565,
246
+ "nll": 1.5063557911660717
247
+ },
248
+ "massive:mn-MN": {
249
+ "count": 200,
250
+ "correct": 100,
251
+ "accuracy": 0.5,
252
+ "nll": 1.578224376673286
253
+ },
254
+ "massive:ms-MY": {
255
+ "count": 200,
256
+ "correct": 158,
257
+ "accuracy": 0.79,
258
+ "nll": 0.7296966203279017
259
+ },
260
+ "massive:my-MM": {
261
+ "count": 200,
262
+ "correct": 100,
263
+ "accuracy": 0.5,
264
+ "nll": 1.6915851091994594
265
+ },
266
+ "massive:nb-NO": {
267
+ "count": 200,
268
+ "correct": 156,
269
+ "accuracy": 0.78,
270
+ "nll": 0.7617846100042545
271
+ },
272
+ "massive:nl-NL": {
273
+ "count": 200,
274
+ "correct": 167,
275
+ "accuracy": 0.835,
276
+ "nll": 0.5973099173809344
277
+ },
278
+ "massive:pl-PL": {
279
+ "count": 200,
280
+ "correct": 154,
281
+ "accuracy": 0.77,
282
+ "nll": 0.6892870971949061
283
+ },
284
+ "massive:pt-PT": {
285
+ "count": 200,
286
+ "correct": 169,
287
+ "accuracy": 0.845,
288
+ "nll": 0.5603875439414431
289
+ },
290
+ "massive:ro-RO": {
291
+ "count": 200,
292
+ "correct": 153,
293
+ "accuracy": 0.765,
294
+ "nll": 0.8143788994751282
295
+ },
296
+ "massive:ru-RU": {
297
+ "count": 200,
298
+ "correct": 169,
299
+ "accuracy": 0.845,
300
+ "nll": 0.4346364665858552
301
+ },
302
+ "massive:sl-SL": {
303
+ "count": 200,
304
+ "correct": 139,
305
+ "accuracy": 0.695,
306
+ "nll": 0.9863898763288126
307
+ },
308
+ "massive:sq-AL": {
309
+ "count": 200,
310
+ "correct": 144,
311
+ "accuracy": 0.72,
312
+ "nll": 1.008332237904707
313
+ },
314
+ "massive:sv-SE": {
315
+ "count": 200,
316
+ "correct": 157,
317
+ "accuracy": 0.785,
318
+ "nll": 0.775370210089804
319
+ },
320
+ "massive:sw-KE": {
321
+ "count": 200,
322
+ "correct": 97,
323
+ "accuracy": 0.485,
324
+ "nll": 1.7013354790948956
325
+ },
326
+ "massive:ta-IN": {
327
+ "count": 200,
328
+ "correct": 131,
329
+ "accuracy": 0.655,
330
+ "nll": 1.1463096412760028
331
+ },
332
+ "massive:te-IN": {
333
+ "count": 200,
334
+ "correct": 124,
335
+ "accuracy": 0.62,
336
+ "nll": 1.2817266688502815
337
+ },
338
+ "massive:th-TH": {
339
+ "count": 200,
340
+ "correct": 147,
341
+ "accuracy": 0.735,
342
+ "nll": 0.8718785998410745
343
+ },
344
+ "massive:tl-PH": {
345
+ "count": 200,
346
+ "correct": 140,
347
+ "accuracy": 0.7,
348
+ "nll": 0.9638123165387646
349
+ },
350
+ "massive:tr-TR": {
351
+ "count": 200,
352
+ "correct": 157,
353
+ "accuracy": 0.785,
354
+ "nll": 0.6703527633025136
355
+ },
356
+ "massive:ur-PK": {
357
+ "count": 200,
358
+ "correct": 131,
359
+ "accuracy": 0.655,
360
+ "nll": 1.1953078716722099
361
+ },
362
+ "massive:vi-VN": {
363
+ "count": 200,
364
+ "correct": 144,
365
+ "accuracy": 0.72,
366
+ "nll": 0.9832935351908162
367
+ },
368
+ "massive:zh-CN": {
369
+ "count": 200,
370
+ "correct": 156,
371
+ "accuracy": 0.78,
372
+ "nll": 0.6844688390764435
373
+ },
374
+ "massive:zh-TW": {
375
+ "count": 200,
376
+ "correct": 154,
377
+ "accuracy": 0.77,
378
+ "nll": 0.7012105530022036
379
+ },
380
+ "open-jev:customer-control-v1": {
381
+ "count": 96,
382
+ "correct": 93,
383
+ "accuracy": 0.96875,
384
+ "nll": 0.11922412518095221
385
+ },
386
+ "open-jev:painting-geometry-v1": {
387
+ "count": 192,
388
+ "correct": 153,
389
+ "accuracy": 0.796875,
390
+ "nll": 0.4879071857538497
391
+ },
392
+ "open-jev:reasoning-control-v1": {
393
+ "count": 131,
394
+ "correct": 68,
395
+ "accuracy": 0.5190839694656488,
396
+ "nll": 1.5929976853600538
397
+ },
398
+ "open-jev:snake-v1": {
399
+ "count": 200,
400
+ "correct": 158,
401
+ "accuracy": 0.79,
402
+ "nll": 0.4846087664592145
403
+ },
404
+ "open-jev:tic_tac_toe-v1": {
405
+ "count": 115,
406
+ "correct": 47,
407
+ "accuracy": 0.40869565217391307,
408
+ "nll": 1.2158676274170932
409
+ },
410
+ "open-jev:tile_platformer-v1": {
411
+ "count": 22,
412
+ "correct": 20,
413
+ "accuracy": 0.9090909090909091,
414
+ "nll": 0.09021685084244535
415
+ },
416
+ "open-jev:trex_runner-v1": {
417
+ "count": 2,
418
+ "correct": 1,
419
+ "accuracy": 0.5,
420
+ "nll": 1.2663163314336559
421
+ },
422
+ "open-jev:vizdoom-basic-v1": {
423
+ "count": 153,
424
+ "correct": 149,
425
+ "accuracy": 0.9738562091503268,
426
+ "nll": 0.03614709415145343
427
+ },
428
+ "open-jev:workflow-controls-v1/agent_trace_observability": {
429
+ "count": 71,
430
+ "correct": 71,
431
+ "accuracy": 1.0,
432
+ "nll": 0.0011360287241133358
433
+ },
434
+ "open-jev:workflow-controls-v1/customer_service": {
435
+ "count": 137,
436
+ "correct": 136,
437
+ "accuracy": 0.9927007299270073,
438
+ "nll": 0.030294906363817344
439
+ },
440
+ "open-jev:workflow-controls-v1/invoice_processing": {
441
+ "count": 112,
442
+ "correct": 110,
443
+ "accuracy": 0.9821428571428571,
444
+ "nll": 0.10505136013108395
445
+ },
446
+ "open-jev:workflow-controls-v1/security_incidents": {
447
+ "count": 121,
448
+ "correct": 121,
449
+ "accuracy": 1.0,
450
+ "nll": 0.009351040906769742
451
+ },
452
+ "replay:agnews": {
453
+ "count": 29,
454
+ "correct": 23,
455
+ "accuracy": 0.7931034482758621,
456
+ "nll": 0.5144041118900763
457
+ },
458
+ "replay:banking77": {
459
+ "count": 52,
460
+ "correct": 50,
461
+ "accuracy": 0.9615384615384616,
462
+ "nll": 0.1474100871587149
463
+ },
464
+ "replay:emotion": {
465
+ "count": 15,
466
+ "correct": 12,
467
+ "accuracy": 0.8,
468
+ "nll": 0.5748594660688843
469
+ },
470
+ "replay:gsm8k": {
471
+ "count": 200,
472
+ "correct": 196,
473
+ "accuracy": 0.98,
474
+ "nll": 0.047936887727627886
475
+ },
476
+ "replay:hellaswag": {
477
+ "count": 119,
478
+ "correct": 46,
479
+ "accuracy": 0.3865546218487395,
480
+ "nll": 3.136717304436293
481
+ },
482
+ "replay:mmlu": {
483
+ "count": 182,
484
+ "correct": 93,
485
+ "accuracy": 0.510989010989011,
486
+ "nll": 2.17472550274021
487
+ },
488
+ "replay:winogrande": {
489
+ "count": 200,
490
+ "correct": 107,
491
+ "accuracy": 0.535,
492
+ "nll": 1.151156462092322
493
+ },
494
+ "typed": {
495
+ "count": 500,
496
+ "correct": 365,
497
+ "accuracy": 0.73,
498
+ "nll": 2.027869497306916
499
+ },
500
+ "typed:choice": {
501
+ "count": 155,
502
+ "correct": 108,
503
+ "accuracy": 0.6967741935483871,
504
+ "nll": 2.5421524021866317
505
+ },
506
+ "typed:noul": {
507
+ "count": 142,
508
+ "correct": 115,
509
+ "accuracy": 0.8098591549295775,
510
+ "nll": 1.5555020525764138
511
+ },
512
+ "typed:score": {
513
+ "count": 203,
514
+ "correct": 142,
515
+ "accuracy": 0.6995073891625616,
516
+ "nll": 1.9656149499935025
517
+ },
518
+ "banking77": {
519
+ "count": 100,
520
+ "correct": 83,
521
+ "accuracy": 0.83,
522
+ "recall_at_16": 0.99,
523
+ "unsupported": 0
524
+ }
525
+ },
526
+ "quality_gate": false
527
+ }
pyproject.toml ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=70"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "supersonic-julia"
7
+ version = "0.1.0"
8
+ requires-python = ">=3.11"
9
+ dependencies = ["torch>=2.6", "transformers>=5.0,<5.1", "safetensors>=0.5", "numpy>=1.26"]
10
+
11
+ [tool.setuptools.packages.find]
12
+ include = ["julia*"]
13
+ exclude = ["julia.router.tests*"]
14
+
15
+ [project.optional-dependencies]
16
+ cuda = ["bitsandbytes>=0.48,<0.50", "accelerate>=1.10,<2"]
17
+
18
+ [tool.setuptools.package-data]
19
+ "julia.router" = ["native/*.bend", "native/*.c", "README.md"]
20
+ "julia" = ["native/*.c", "native/*.bend", "native/*.json"]
tokenizer/tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:609d8f4c067cd3950f88594c5a802616cea245823836ef5848ee4fc40aab5b6f
3
+ size 34363188
tokenizer/tokenizer_config.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "bos_token": "<bos>",
4
+ "clean_up_tokenization_spaces": false,
5
+ "cls_token": "<bos>",
6
+ "eos_token": "<eos>",
7
+ "extra_special_tokens": [
8
+ "<start_of_turn>",
9
+ "<end_of_turn>"
10
+ ],
11
+ "is_local": true,
12
+ "mask_token": "<mask>",
13
+ "max_length": 512,
14
+ "model_input_names": [
15
+ "input_ids",
16
+ "attention_mask"
17
+ ],
18
+ "model_max_length": 8192,
19
+ "pad_token": "<pad>",
20
+ "padding_side": "right",
21
+ "sep_token": "<eos>",
22
+ "spaces_between_special_tokens": false,
23
+ "stride": 0,
24
+ "tokenizer_class": "TokenizersBackend",
25
+ "truncation_side": "right",
26
+ "truncation_strategy": "longest_first",
27
+ "unk_token": "<unk>"
28
+ }