kleeedolinux commited on
Commit ·
5278d6b
0
Parent(s):
Publish Julia 1 model and Python runtime
Browse files- .gitattributes +39 -0
- .gitignore +9 -0
- README.md +138 -0
- assets/julia-banner.png +3 -0
- assets/supersonic-labs-dark.svg +3 -0
- assets/supersonic-labs-light.svg +3 -0
- encoder/config.json +79 -0
- inference-policy.json +10 -0
- julia/__init__.py +8 -0
- julia/cuda.py +56 -0
- julia/data.py +142 -0
- julia/inference.py +74 -0
- julia/model.py +133 -0
- julia/probabilities.py +17 -0
- julia/router/.gitignore +3 -0
- julia/router/README.md +157 -0
- julia/router/__init__.py +12 -0
- julia/router/build.py +62 -0
- julia/router/encoder.py +32 -0
- julia/router/engine.py +225 -0
- julia/router/native.py +159 -0
- julia/router/native/LAWS.bend +57 -0
- julia/router/native/PROOF.bend +50 -0
- julia/router/native/bridge.c +452 -0
- julia/router/native/router.bend +360 -0
- julia/router/router.py +173 -0
- julia/router/specialize.py +38 -0
- julia/router/tests/synthetic.py +36 -0
- julia/router/tests/test_engine.py +367 -0
- julia/router/tests/test_router.py +108 -0
- julia/router/transformer.py +97 -0
- julia_config.json +8 -0
- metrics/accuracy-20260924.json +434 -0
- metrics/validation.json +45 -0
- model.safetensors +3 -0
- provenance.json +527 -0
- pyproject.toml +20 -0
- tokenizer/tokenizer.json +3 -0
- tokenizer/tokenizer_config.json +28 -0
.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 |
+

|
| 13 |
+
|
| 14 |
+
<h1 align="center">Julia 1</h1>
|
| 15 |
+
<p align="center">From context to decisions.</p>
|
| 16 |
+
|
| 17 |
+

|
| 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
|
assets/supersonic-labs-dark.svg
ADDED
|
|
Git LFS Details
|
assets/supersonic-labs-light.svg
ADDED
|
|
Git LFS Details
|
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 |
+
}
|