diff --git a/Cargo.lock b/Cargo.lock index 2bb95d2..c43ff22 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -155,6 +155,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -236,6 +245,12 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -304,6 +319,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + [[package]] name = "daachorse" version = "3.0.3" @@ -391,8 +415,19 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", ] [[package]] @@ -683,6 +718,15 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "hybrid-array" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27f864f10dfb56725ce5ce5472bc52252c8f93a4ab86327122cebf62c5f59a17" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.11.1" @@ -1092,7 +1136,7 @@ dependencies = [ "safetensors 0.6.2", "serde", "serde_json", - "sha2", + "sha2 0.10.9", ] [[package]] @@ -1110,7 +1154,7 @@ dependencies = [ "safetensors 0.8.0", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "tokenizers 0.22.2", "tokio", ] @@ -1133,6 +1177,23 @@ dependencies = [ "tokio", ] +[[package]] +name = "omni-jev-vl-native" +version = "0.1.0" +dependencies = [ + "anyhow", + "axum", + "half", + "omni-qwen3-5-native", + "omni-runtime", + "safetensors 0.8.0", + "serde_json", + "sha2 0.11.0", + "tempfile", + "tokenizers 0.22.2", + "tokio", +] + [[package]] name = "omni-laya" version = "0.1.0" @@ -1146,7 +1207,7 @@ dependencies = [ "safetensors 0.6.2", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "tempfile", "tokenizers 0.23.2", "tokio", @@ -1702,7 +1763,18 @@ checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", "cpufeatures 0.2.17", - "digest", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.1", + "digest 0.11.3", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 242d51c..dc0bef3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/laya", "src/backends/cuda"] +members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/jev_vl/native", "src/models/laya", "src/backends/cuda"] resolver = "3" diff --git a/README.md b/README.md index cbcb633..e940c6a 100644 --- a/README.md +++ b/README.md @@ -167,6 +167,7 @@ and CLM has a stub-encoder contract recipe: | Cua-S1 4B 0.2 (`text` adapter) | [Python worker](recipe/cua_s1/text.md); [native worker](recipe/cua_s1/native.md), CUDA, run on sm_89 | | Cua-S1 4B 0.2 (`multimodal` adapter) | [Python CUDA worker](src/frontend/cua_s1.py); one PNG/JPEG screenshot, `choice`; native screenshot execution remains in progress | | Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](recipe/open_jev/native.md); eager independent text candidates; [H200 validation](recipe/open_jev/validation.md) | +| autotrust/JEV-27B-VL | [Experimental Rust/CUDA worker](recipe/jev_vl/README.md); single-question text and offline-preencoded image inputs; [bounded H800 validation and limits](recipe/jev_vl/validation.md) | | CLM-v0.1-8B | [External worker with a CPU stub encoder](recipe/clm/README.md); contract checks only, real Qwen3-8B decisions unverified by this recipe | [Supported models and hardware](docs/supported-models.md) lists the devices diff --git a/docs/architecture.md b/docs/architecture.md index e4612f7..badc4c1 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -59,6 +59,13 @@ the existing FP32 rounding and bias order. Finishing checks output cardinality before reconstruction. HTTP validation, error status/body conventions, and real warmup before readiness remain model-specific and unchanged. +The experimental [JEV-VL worker](../recipe/jev_vl/README.md) also uses the shared +Qwen executor. It admits one official single-decision request at a time, reads +selected LM-head rows, and can retain image-prefix state across requests. Image +features must be prepared offline. Its `{kind, state, question, options}` contract +is distinct from the existing `{model, state, questions}` envelope; frontend +transport alone does not adapt those schemas. + ## Layer ownership and implementation language | Component | Owns | Native target implementation | diff --git a/docs/supported-models.md b/docs/supported-models.md index f9a8539..40d40b3 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -1,6 +1,7 @@ # Supported models and hardware -This page covers what runs from `main`. Start with the +This page distinguishes validated serving paths from experimental integrations. +Start with the [CPU decision walkthrough](getting-started.md). The [rolling model tracker (#83)](https://github.com/ThinkFlowLab/system1-omni/issues/83) records available paths separately from proposed integrations and hardware evidence. @@ -15,6 +16,7 @@ Models that are being added are also tracked in issues labeled [new model](https | Cua-S1 4B 0.2, `text` adapter | [Native Rust worker](../recipe/cua_s1/native.md) on the [Qwen3.5 CUDA kernels](../src/backends/cuda/qwen3_5/README.md) | Not supported | Validated on compute capability 8.9 ([#19](https://github.com/ThinkFlowLab/system1-omni/pull/19), [#52](https://github.com/ThinkFlowLab/system1-omni/pull/52)) | Not supported | Compute capability 8.0 or newer, the CUDA toolkit to build, weights merged with `export_text_merged.py` | | Cua-S1 4B 0.2, `multimodal` adapter | Reference worker on Transformers and PEFT, [`src/frontend/cua_s1.py`](../src/frontend/cua_s1.py); no recipe yet | Not supported | Validated ([#17](https://github.com/ThinkFlowLab/system1-omni/pull/17), [#18](https://github.com/ThinkFlowLab/system1-omni/pull/18)) | Not supported | The state is one PNG or JPEG image; upstream's `weights.lock.json` next to the base weights | | Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](../recipe/open_jev/native.md) on the shared Qwen3.5/3.8 executor | Not supported | Validated on H200 (sm_90) for the [74 single-candidate workload](../recipe/open_jev/validation.md) | Not supported | Compute capability 8.0 or newer, CUDA toolkit to build, exported merged weights and trained head | +| autotrust/JEV-27B-VL | [Experimental native worker](../recipe/jev_vl/README.md), single-question API | Contract tests only; no CPU inference | H800 frozen-corpus check; [experimental scope and limits](../recipe/jev_vl/validation.md) | Not supported | Pinned merged checkpoint, rebuilt ABI 6 CUDA library; image inputs require offline Transformers preencoding | | CLM-v0.1-8B | [External `clm-serve` recipe](../recipe/clm/README.md) with a CPU stub embeddings server | **Stub-encoder contract checks only** ([#23](https://github.com/ThinkFlowLab/system1-omni/pull/23)); not real Qwen3-8B decisions | Real encoder unverified by the merged recipe | Unverified | Python, upstream CLM and head checkpoint; a real encoder requires a separate embeddings server | - **Validated:** covered by the recipe on `main` or by the checks in the linked merged pull request. diff --git a/mkdocs.yml b/mkdocs.yml index 39e0d74..f6d317b 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -67,5 +67,7 @@ nav: - Cua-S1 native text worker: recipe/cua_s1/native.md - Open-Jev native text worker: recipe/open_jev/native.md - Open-Jev H200 validation: recipe/open_jev/validation.md + - JEV-27B-VL experimental worker: recipe/jev_vl/README.md + - JEV-27B-VL validation status: recipe/jev_vl/validation.md - CLM stub-encoder contract: recipe/clm/README.md - Contributing: CONTRIBUTING.md diff --git a/recipe/README.md b/recipe/README.md index 7d64dcf..39e3e6b 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -15,6 +15,8 @@ For a first real decision, follow the [complete CPU walkthrough](../docs/getting the Rust worker, export the merged weights and start the worker. - [Open-Jev-27B-v1.1 native text worker](open_jev/native.md): export the merged text backbone and trained decision head, then serve with Rust and CUDA. +- [JEV-27B-VL experimental worker](jev_vl/README.md): native text and preencoded + image decisions; bounded H800 validation with explicit deployment limits. - [CLM behind the frontend](clm/README.md): run CLM's own server behind the frontend on CPU with a stub encoder, and what the response comparison has to allow for. diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 0e35200..83322b2 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -28,7 +28,7 @@ each exact prompt length warms the GEMM plans and captures the forward pass; later requests replay it with freshly uploaded token ids. At most eight lengths are cached. Growing the scratch allocation clears the captures before freeing their buffers. Capture adds first-use latency; leave the variable unset to use -the eager control. Rebuild both the worker and CUDA library together (ABI 4). +the eager control. Rebuild both the worker and CUDA library together (ABI 6). If capture fails, the worker returns the completed eager result and disables Graph capture/replay for its remaining lifetime, logging the failure to stderr. diff --git a/recipe/jev_vl/README.md b/recipe/jev_vl/README.md new file mode 100644 index 0000000..df8109a --- /dev/null +++ b/recipe/jev_vl/README.md @@ -0,0 +1,158 @@ +# JEV-27B-VL experimental native recipe + +This recipe serves `autotrust/JEV-27B-VL` System-1 decisions with Rust/CUDA. +Text runs natively; images must first be encoded offline by Transformers. +The [model contract](../../src/models/jev_vl/README.md) explains the single-question +API, verbalizer head and cache ownership. The reviewed candidate passed a +[bounded H800 validation](validation.md#historical-h800-validation); +this remains an experimental integration with the documented coverage limits. + +## Prepare a pinned checkpoint + +Run all commands from the repository root. Preparation used Linux, Python 3.10, +PyTorch 2.13.0+cu130, Transformers 5.17.0 and safetensors 0.8.0. The +[requirements file](requirements.txt) pins observed direct dependencies; a clean +installation of those pins has not been revalidated. Use an isolated environment: + +```sh +python3.10 -m venv .venv-jev-vl +.venv-jev-vl/bin/python -m pip install -r recipe/jev_vl/requirements.txt +.venv-jev-vl/bin/hf download autotrust/JEV-27B-VL \ + --revision f34b598d4ef4bcefd337bee8d8e7ddd3b7733ccc \ + --local-dir weights/JEV-27B-VL +CUDA_VISIBLE_DEVICES='' .venv-jev-vl/bin/python recipe/jev_vl/export_merged.py \ + --model weights/JEV-27B-VL --out weights/jev-vl-merged --max-length 16384 +``` + +The CPU exporter streams shards, merges the backbone LoRA in FP32 before BF16 +rounding, and separately exports selected merged LM-head rows as FP32. Keep the +original source checkpoint for image preprocessing. The historical language +export occupied about 48 GiB in addition to the source checkpoint. Peak host +RAM was not measured; the export job requested 64 GiB. Do not assume the whole +pipeline fits on a low-memory workstation from the streaming implementation +alone. The output directory must not exist before export. + +## Build and launch + +The author measured the worker on one H800 80 GB (`sm_90`). A maintainer also +[reported a bounded L20X replay](https://github.com/ThinkFlowLab/system1-omni/pull/96#issuecomment-6018324526) +on an earlier revision. These runs do not validate later code changes or maximum +context lengths. Use an allocated GPU on +scheduled hosts. Build the CUDA library and Rust workers from the same revision; +the integrated backend uses **ABI 6** and older libraries must be rebuilt. + +```sh +src/backends/cuda/qwen3_5/build.sh target/release 90 +cargo build --release --locked -p omni-jev-vl-native -p omni-jev +JEV_VL_MODEL=weights/jev-vl-merged \ + JEV_VL_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + JEV_VL_CACHE=0 target/release/omni-jev-vl-native +``` + +`JEV_VL_HOST` defaults to `127.0.0.1`, `JEV_VL_PORT` to `8001`, and the CUDA +library defaults to the file beside the executable. The worker uses visible +CUDA device 0. `/health` becomes available after model loading and a successful +text warmup; that health check does not validate image assets. + +In separate terminals, start the frontend and send the same request directly and +through it: + +```sh +OMNI_JEV_BIND=127.0.0.1:8080 OMNI_JEV_BACKEND_URL=http://127.0.0.1:8001 \ + target/release/omni-jev +``` + +```sh +curl --fail-with-body http://127.0.0.1:8001/health +curl --fail-with-body http://127.0.0.1:8080/health +curl --fail-with-body http://127.0.0.1:8001/v1/systemone \ + -H 'Content-Type: application/json' --data-binary @recipe/jev_vl/example-request.json +curl --fail-with-body http://127.0.0.1:8080/v1/systemone \ + -H 'Content-Type: application/json' --data-binary @recipe/jev_vl/example-request.json +``` + +Expect a 200 response with two finite `probabilities`, their sum approximately +one, an in-range `choice_index`, and the corresponding string `choice`. Compare +decision fields and usage between both responses; `elapsed_seconds` varies. +This worker takes `kind`, `state`, `question` and `options`, rather than the +`questions`/`answers` envelope in the generic comparison recipe. A fixed expected +choice is not asserted here; use the frozen comparison corpus for numerical +checks. The frontend routes `/health` and `/v1/systemone`; it does not proxy +generation routes such as `/v1/chat/completions`. An unknown route returns the +frontend's own 404 rather than this worker's error envelope. + +## Prepare images before serving + +The worker does not decode images, fetch URLs or invoke a vision tower on cache +miss. First prepare a JSONL manifest with one object per line containing a +`request` whose `state` includes `{"image":"data:image/png;base64,..."}`. The +preencoder accepts base64 data URIs; a fresh URL or changed image needs a new +asset. It runs on a GPU and should finish before starting the language worker +on the same device: + +```sh +.venv-jev-vl/bin/python recipe/jev_vl/preencode.py \ + --model weights/JEV-27B-VL --manifest /path/to/requests.jsonl \ + --out weights/jev-vl-image-assets +JEV_VL_MODEL=weights/jev-vl-merged JEV_VL_IMGCACHE=weights/jev-vl-image-assets \ + JEV_VL_CACHE=1 target/release/omni-jev-vl-native +``` + +Replace the manifest path with your own prepared workload. The preencoder writes +`sha256(exact_data_uri)/emb.safetensors` and `grid.json`; the worker requires the +exact same URI string. Use a new output directory for a changed checkpoint or +processor. Complete assets are skipped; interrupted writes are regenerated on +the next run. Source identity is not a cryptographic runtime compatibility check. Preencoding, +including image decoding and vision execution, is excluded from the reported +worker timing. This is useful for repeated decisions on a prepared image; it is +not an end-to-end live screenshot service. + +## Cache controls + +| Variable | Default | Meaning | +| --- | --- | --- | +| `JEV_VL_CACHE` | `1` | Master enable; `0` uses a full language forward and rereads prepared image assets. | +| `JEV_VL_L1`, `JEV_VL_L2`, `JEV_VL_L3` | `1` | Processor records, parsed image assets, and language-prefix state. | +| `JEV_VL_L1_MAX` | `256` | Maximum resident processor records. | +| `JEV_VL_L2_BYTES` | `1073741824` | Parsed image-asset cache budget. | +| `JEV_VL_L3_BYTES` | `2147483648` | Cached language-prefix device-state budget. | + +Cache budgets do not include weights, model scratch, in-flight request inputs or +total process memory. `GET /v1/cache/stats` exposes counters and resident cache +accounting. Use those counters to assess L2 hits; `x-jev-cache` does not report +an L2 hit field. `POST /v1/cache/reset` resets counters only; restart the worker for +a cold cache. Request preparation may overlap, while GPU execution stays serial. +Use cache modes as experimental controls within the documented validation scope. + +## Local checks + +CPU checks do not need weights or CUDA: + +```sh +cargo fmt --all --check +cargo clippy --workspace --locked --all-targets -- -D warnings +cargo test --workspace --locked +cargo build --workspace --release --locked +python3 -m unittest discover -s tests/benchmarks -p 'test_jev_vl_*.py' -v +python3 recipe/jev_vl/preencode.py --help +``` + +The image-prefix tokenizer check requires the exported checkpoint. First +[download and restore the frozen corpus](validation.md#download-the-frozen-corpus). +It uses the corpus's fixed `[1, 60, 60]` grid and synthetic embedding +rows, so it checks token/position splitting rather than the vision encoder. +CUDA kernel tests require an allocated GPU and rebuilt ABI 6 library: + +```sh +JEV_VL_EXPORT=$PWD/weights/jev-vl-merged \ + JEV_VL_MANIFEST="$jev_vl_evidence/manifest.jsonl" \ + cargo test --locked -p omni-jev-vl-native --test jev_vl_prefix -- --ignored +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + QWEN3_5_CHECKPOINT=$PWD/weights/jev-vl-merged \ + cargo test --release --locked -p omni-qwen3-5-native --test kernels -- \ + --ignored --test-threads=1 +``` + +These checks do not replace full-checkpoint parity or direct/frontend HTTP +validation. The [validation page](validation.md) lists the remaining gates and +separates historical measurements from the current revision. diff --git a/recipe/jev_vl/example-request.json b/recipe/jev_vl/example-request.json new file mode 100644 index 0000000..ce62504 --- /dev/null +++ b/recipe/jev_vl/example-request.json @@ -0,0 +1,7 @@ +{ + "model": "autotrust/JEV-27B-VL", + "kind": "choice", + "state": "The customer received a package with a broken screen and asks for help.", + "question": "What should support do next?", + "options": ["Offer a replacement", "Close the ticket without a reply"] +} diff --git a/recipe/jev_vl/export_merged.py b/recipe/jev_vl/export_merged.py new file mode 100644 index 0000000..36f99e9 --- /dev/null +++ b/recipe/jev_vl/export_merged.py @@ -0,0 +1,201 @@ +"""Export the merged JEV-27B-VL text backbone plus the verbalizer label head, on CPU. + +Streaming merge, low RAM: the script never materializes the whole model. It walks +the base checkpoint shard by shard, adds the trained vLLM-LoRA deltas +(``scale * (lora_B @ lora_A)``, ``scale = lora_alpha / r``) computed in float32 and +rounded once back to bfloat16, and rewrites only the ``model.language_model.*`` +tensors under their original names, so the native backend's prefix autodetection +loads the output unchanged. ``lm_head.weight`` is merged the same way and the rows +the System-1 verbalizer readout needs (the union of the trained 24 head slots and +the tokenizer's single-token option labels) land in ``label_head.safetensors`` +(float32). ``jev_vl_export.json`` — written last — carries the decision semantics: +per-kind temperatures, the 24-slot bias/ranges, the exported label list, and the +checkpoint pins. System-1 uses a raw prompt without a chat template. + +See recipe/jev_vl/README.md for the pinned export environment. +No PEFT and no GPU are needed. +""" + +import argparse +import hashlib +import json +import shutil +import string +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoTokenizer + +EXPORT_FORMAT = "jev-vl-text-merged/1" +MODEL_ID = "autotrust/JEV-27B-VL" +PROTOCOL = "jev27-bare-v1" +MAX_OPTIONS = 256 +TOKENIZER_FILES = [ + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "special_tokens_map.json", + "added_tokens.json", + "chat_template.jinja", +] + + +def sha256(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 22), b""): + h.update(chunk) + return h.hexdigest() + + +def hf_revision(model: Path) -> str | None: + meta = model / ".cache/huggingface/download/config.json.metadata" + try: + return meta.read_text().splitlines()[0].strip() + except OSError: + return None + + +def load_adapter(adapter: Path) -> tuple[dict[str, tuple[torch.Tensor, torch.Tensor]], float]: + cfg = json.loads((adapter / "adapter_config.json").read_text()) + scale = cfg["lora_alpha"] / cfg["r"] + assert cfg["peft_type"] == "LORA" and cfg["inference_mode"] and cfg["lora_dropout"] == 0.0 + pairs: dict[str, dict[str, torch.Tensor]] = {} + with safe_open(str(adapter / "adapter_model.safetensors"), framework="pt") as f: + for key in f.keys(): + assert key.startswith("base_model.model.") and key.endswith(".weight"), key + base = key[len("base_model.model."):] + name, kind = base.rsplit(".lora_", 1) + name += ".weight" + pairs.setdefault(name, {})[kind.split(".")[0].lower()] = f.get_tensor(key) + out = {} + for name, ab in pairs.items(): + assert set(ab) == {"a", "b"}, f"{name}: adapter pair incomplete" + a, b = ab["a"].float(), ab["b"].float() + assert a.shape[0] == cfg["r"] and b.shape[1] == cfg["r"], f"{name}: rank mismatch" + out[name] = (a, b) + return out, scale + + +def single_token_labels(tok, context: str) -> list[tuple[str, int]]: + """serve_decide.py's option-label scan, verbatim.""" + out = [] + for lab in list(string.ascii_uppercase) + [a + b for a in string.ascii_uppercase + for b in string.ascii_uppercase]: + t = tok.encode(lab, add_special_tokens=False) + if len(t) == 1 and t[0] in tok.encode(context.format(lab), add_special_tokens=False): + out.append((lab, t[0])) + if len(out) == MAX_OPTIONS: + break + return out + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--model", required=True, type=Path) + parser.add_argument("--out", required=True, type=Path) + parser.add_argument("--max-length", type=int, default=16384) + args = parser.parse_args() + model, out = args.model, args.out + if out.exists(): + raise ValueError("output already exists; choose a new export directory") + if not 1 <= args.max_length <= 32768: + raise ValueError("max length must be within 1..=32768") + adapter = model / "adapter_vllm" + head_cfg = json.loads((adapter / "decision_head.json").read_text()) + temperatures = json.loads((model / "calibration.json").read_text())["per_kind"] + assert sorted(head_cfg["slots"]["ranges"]) == ["choice", "noul", "score"] + assert head_cfg["slots"]["ranges"]["noul"] == [0, 2] + assert head_cfg["slots"]["ranges"]["choice"] == [8, 24] + assert len(head_cfg["verbalizer_ids"]) == len(head_cfg["bias"]) == 24 + deltas, scale = load_adapter(adapter) + assert "lm_head.weight" in deltas, "adapter must train lm_head" + lm_delta = deltas.pop("lm_head.weight") + + tokenizer = AutoTokenizer.from_pretrained(model, local_files_only=True) + labels = single_token_labels(tokenizer, "x\n{}) y") + assert len(labels) == MAX_OPTIONS, f"only {len(labels)} single-token labels" + lo, hi = head_cfg["slots"]["ranges"]["choice"] + assert [t for _, t in labels[: hi - lo]] == head_cfg["verbalizer_ids"][lo:hi], \ + "first labels must be the trained A-P head" + index = json.loads((model / "model.safetensors.index.json").read_text()) + weight_map: dict[str, str] = index["weight_map"] + files = sorted(set(weight_map.values())) + out.mkdir(parents=True) + consumed: set[str] = set() + new_map: dict[str, list[str]] = {} + print(f"merging with LoRA scale={scale} from {adapter}", flush=True) + for shard in files: + merged: dict[str, torch.Tensor] = {} + with safe_open(str(model / shard), framework="pt") as f: + for name in sorted(f.keys()): + if not name.startswith("model.language_model."): + continue + w = f.get_tensor(name) + if name in deltas: + a, b = deltas[name] + w = (w.float() + scale * (b @ a)).to(torch.bfloat16) + consumed.add(name) + merged[name] = w.contiguous() + if merged: + save_file(merged, str(out / shard)) + for name in merged: + new_map.setdefault(shard, []).append(name) + print(f" {shard}: {len(merged)} language tensors", flush=True) + missing = set(deltas) - consumed + assert not missing, f"adapter tensors without a base tensor: {sorted(missing)[:4]}" + (out / "model.safetensors.index.json").write_text(json.dumps({ + "metadata": {"total_size": 0, "jev_vl_export": EXPORT_FORMAT}, + "weight_map": {name: shard for shard, names in sorted(new_map.items()) for name in sorted(names)}, + }) + "\n") + + # The merged label rows for the verbalizer readout (float32). + head_ids = list(dict.fromkeys(head_cfg["verbalizer_ids"] + [t for _, t in labels])) + lm_shard = weight_map["lm_head.weight"] + with safe_open(str(model / lm_shard), framework="pt") as f: + lm_head = f.get_tensor("lm_head.weight") + a, b = lm_delta + idx = torch.tensor(head_ids) + rows = lm_head[idx].float() + scale * (b[idx] @ a) + assert torch.isfinite(rows).all(), "non-finite merged label rows" + save_file({"rows": rows.contiguous(), + "ids": idx.to(torch.int64)}, str(out / "label_head.safetensors")) + + for name in TOKENIZER_FILES: + src = model / name + if src.exists(): + shutil.copy2(src, out / name) + shutil.copy2(model / "config.json", out / "config.json") + + manifest = { + "format": EXPORT_FORMAT, + "model_id": MODEL_ID, + "protocol": PROTOCOL, + "hf_revision": hf_revision(model), + "pins": { + "model_index_sha256": sha256(model / "model.safetensors.index.json"), + "adapter_config_sha256": sha256(adapter / "adapter_config.json"), + "adapter_model_sha256": sha256(adapter / "adapter_model.safetensors"), + "decision_head_sha256": sha256(adapter / "decision_head.json"), + "calibration_sha256": sha256(model / "calibration.json"), + }, + "merge": {"lora_scale": scale, "delta_dtype": "float32", "backbone_dtype": "bfloat16"}, + "temperatures": temperatures, + "verbalizer_ids": head_cfg["verbalizer_ids"], + "verbalizer_bias": head_cfg["bias"], + "slots": head_cfg["slots"], + "labels": [l for l, _ in labels], + "label_ids": [t for _, t in labels], + "label_head": {"file": "label_head.safetensors", "ids": head_ids, "dtype": "float32"}, + "max_length": args.max_length, + } + # Written last: the native worker refuses incomplete exports or plain base weights. + (out / "jev_vl_export.json").write_text(json.dumps(manifest, ensure_ascii=False) + "\n") + print(f"exported {len(new_map)} shards + {len(head_ids)} label rows -> {out}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/recipe/jev_vl/preencode.py b/recipe/jev_vl/preencode.py new file mode 100644 index 0000000..807f98b --- /dev/null +++ b/recipe/jev_vl/preencode.py @@ -0,0 +1,129 @@ +"""Pre-encode manifest images into the (url-hash keyed) adapted-embedding assets. + +Runs the official HF vision tower (same model the vLLM reference path used) on the +GPU once per unique image in a manifest, and stores per-image +``//{emb.safetensors,grid.json}`` with the adapted rows +(``model.model.visual(pixel_values, grid_thw).pooler_output``, shape +[n = prod(grid)/merge^2, 5120], bfloat16) plus the patch grid. The L2 cache holds +these assets; the native worker never decodes images. +""" +import argparse +import hashlib +import io +import json +from pathlib import Path + + +def load_vision(model_dir: Path) -> "torch.nn.Module": + """Hand-assemble the vision tower: no accelerate, no full-model device map.""" + import safetensors + import torch + from transformers import AutoConfig + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel + + cfg = AutoConfig.from_pretrained(model_dir, local_files_only=True) + vis = Qwen3_5VisionModel(cfg.vision_config) + index = json.loads((model_dir / "model.safetensors.index.json").read_text())["weight_map"] + names = {n for n in index if n.startswith("model.visual.")} + tensors = {} + for f in sorted({index[n] for n in names}): + with safetensors.safe_open(str(model_dir / f), framework="pt") as h: + for n in h.keys(): + if n in names: + tensors[n[len("model.visual."):]] = h.get_tensor(n) + missing, unexpected = vis.load_state_dict(tensors, strict=True) + assert not missing and not unexpected, (missing, unexpected) + return vis.to(torch.bfloat16).to("cuda").eval() + + +def images_of(manifest: Path) -> list[str]: + urls = [] + for line in manifest.read_text().splitlines(): + if not line.strip(): + continue + r = json.loads(line)["request"] + state = r.get("state", "") + if isinstance(state, list): + for p in state: + if not isinstance(p, dict): + continue + # Match contract::parts: untyped image_url fields are text, + # and the image shorthand takes precedence over typed parts. + if "image" in p: + url = p["image"] + elif p.get("type") == "image_url": + image_url = p.get("image_url") + url = image_url.get("url") if isinstance(image_url, dict) else None + else: + continue + if not isinstance(url, str) or not url: + raise ValueError("image part must contain a nonempty URL string") + urls.append(url) + return sorted(set(urls)) + + +def decode_url(url: str): + from PIL import Image + + if not url.startswith("data:"): + raise ValueError("recipe supports data: URIs (manifest uses them)") + header, payload = url.split(",", 1) + if ";base64" not in header: + raise ValueError("data URI without base64 encoding") + import base64 + + return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB") + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--model", required=True, type=Path) + ap.add_argument("--manifest", required=True, type=Path) + ap.add_argument("--out", required=True, type=Path) + args = ap.parse_args() + urls = images_of(args.manifest) + print(f"{len(urls)} unique images", flush=True) + if not urls: + return + import torch + from safetensors.torch import save_file + from transformers import AutoProcessor + + model_dir = args.model + processor = AutoProcessor.from_pretrained(model_dir, local_files_only=True) + model = load_vision(model_dir) + args.out.mkdir(parents=True, exist_ok=True) + for url in urls: + key = hashlib.sha256(url.encode()).hexdigest() + dest = args.out / key + if (dest / "emb.safetensors").is_file() and (dest / "grid.json").is_file(): + print(f"{key}: cached", flush=True) + continue + im = decode_url(url) + inputs = processor(images=[im], return_tensors="pt") + grid = inputs["image_grid_thw"][0].tolist() + pixel_values = inputs["pixel_values"].to("cuda", dtype=torch.bfloat16) + grid_thw = torch.tensor([grid], device="cuda") + with torch.no_grad(): + emb = model(pixel_values, grid_thw=grid_thw).pooler_output + emb = emb.detach().cpu().to(torch.bfloat16) + n = torch.prod(torch.tensor(grid)).item() // model.spatial_merge_size**2 + assert emb.shape == (n, 5120), f"{key}: {emb.shape} vs n={n}" + assert torch.isfinite(emb.float()).all(), f"{key}: non-finite" + dest.mkdir(parents=True, exist_ok=True) + # Publish metadata last so an interrupted write is retried on the next run. + (dest / "grid.json").unlink(missing_ok=True) + save_file({"rows": emb.contiguous()}, str(dest / "emb.safetensors")) + grid_tmp = dest / "grid.json.tmp" + grid_tmp.write_text(json.dumps({ + "url_sha256": key, "grid_thw": grid, "n_tokens": n, "dtype": "bfloat16", + "shape": [int(emb.shape[0]), int(emb.shape[1])], + "model_index_sha256": hashlib.sha256((model_dir / "model.safetensors.index.json") + .read_bytes()).hexdigest(), + }) + "\n") + grid_tmp.replace(dest / "grid.json") + print(f"{key}: grid={grid} n={n} rows={tuple(emb.shape)}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/recipe/jev_vl/replay.py b/recipe/jev_vl/replay.py new file mode 100644 index 0000000..380db1b --- /dev/null +++ b/recipe/jev_vl/replay.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Serial JEV-VL replay: HTTP + JSON decode timing, fixed 0.025 probability gate. + +Request serialization and reference comparison are outside client_s. Warmup is +repeated before every pass, recorded, and excluded from the latency summary. +This client does not start a server, reset caches, or establish GPU execution. +""" + +import argparse +import hashlib +import http.client +import json +import math +from pathlib import Path +import re +import statistics +import sys +import time +import urllib.error +import urllib.request + +TOLERANCE = 0.025 +FIELDS = ("kind", "effective_kind", "options", "choice_index", "choice") + + +def invalid_constant(value): + raise ValueError(f"non-finite JSON constant: {value}") + + +def finite_float(value): + parsed = float(value) + return parsed if math.isfinite(parsed) else invalid_constant(value) + + +def loads(raw): + return json.loads(raw, parse_constant=invalid_constant, parse_float=finite_float) + + +def probabilities(body): + values, options, index = body["probabilities"], body["options"], body["choice_index"] + if not isinstance(options, list) or not isinstance(values, list) or not values or len(values) != len(options): + raise ValueError("probability/option length mismatch") + if any(type(p) not in (int, float) or not 0 <= p <= 1 or not math.isfinite(p) for p in values): + raise ValueError("invalid probability") + if abs(sum(values) - 1) > 1e-4: + raise ValueError("probabilities do not sum to one") + if type(index) is not int or index != max(range(len(values)), key=values.__getitem__): + raise ValueError("choice_index differs from probability argmax") + if body["choice"] != options[index]: + raise ValueError("choice differs from the selected option") + return values + + +class NoRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, *args): + return None + + +def post(opener, endpoint, payload): + data = json.dumps(payload, ensure_ascii=False, allow_nan=False).encode() + request = urllib.request.Request(endpoint, data=data, headers={"Content-Type": "application/json"}) + row = {"status": None, "ok": False, "max_diff": None} + start = time.monotonic() + try: + try: + response = opener.open(request, timeout=300) + except urllib.error.HTTPError as error: + response = error + with response: + row["status"] = response.status + row["raw_body"] = response.read().decode("utf-8") + row["body"] = loads(row["raw_body"]) + except (OSError, ValueError, urllib.error.URLError, http.client.HTTPException) as error: + row["error"] = str(error) + row["client_s"] = time.monotonic() - start + return row + + +def compare(row, reference): + try: + if row["status"] != 200 or "error" in row: + raise ValueError(row.get("error", f"HTTP {row['status']}")) + body = row["body"] + current, expected = probabilities(body), probabilities(reference) + row["probabilities"] = current + if len(current) != len(expected): + raise ValueError("reference probability length mismatch") + row["max_diff"] = max(abs(a - b) for a, b in zip(current, expected)) + if any(body[field] != reference[field] for field in FIELDS): + raise ValueError("decision mismatch") + if row["max_diff"] > TOLERANCE: + raise ValueError("probability tolerance exceeded") + row["ok"] = True + except (KeyError, TypeError, ValueError, IndexError) as error: + row["error"] = str(error) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + for name in ("base", "manifest", "reference", "out"): + parser.add_argument(f"--{name}", required=True) + parser.add_argument("--pattern", default="", help="manifest ID prefix (default: all)") + parser.add_argument("--warmup", type=int, default=4, help="requests before each pass") + parser.add_argument("--passes", type=int, default=2) + args = parser.parse_args() + if args.warmup < 0 or args.passes < 1: + parser.error("warmup must be nonnegative and passes must be positive") + raw = Path(args.manifest).read_bytes() + entries = [loads(line) for line in raw.splitlines() if line.strip()] + ids = [entry["id"] for entry in entries] + if not entries or len(ids) != len(set(ids)) or any(not re.fullmatch(r"[\w.-]+", key) for key in ids): + raise ValueError("manifest must have nonempty, unique, safe IDs") + entries = [entry for entry in entries if entry["id"].startswith(args.pattern)] + if not entries: + raise ValueError("pattern selects no requests") + references, hashes = {}, {} + for entry in entries: + if not isinstance(entry["request"], dict): + raise ValueError("request must be an object") + content = (Path(args.reference) / f"{entry['id']}.json").read_bytes() + ref = loads(content) + if ref["id"] != entry["id"] or ref["status"] != 200: + raise ValueError("reference must match the ID and have status 200") + probabilities(ref["body"]) + references[entry["id"]] = ref["body"] + hashes[entry["id"]] = hashlib.sha256(content).hexdigest() + out = Path(args.out) + out.mkdir(parents=True, exist_ok=False) + config = vars(args) | {"tolerance": TOLERANCE, "manifest_sha256": hashlib.sha256(raw).hexdigest(), + "reference_sha256": hashes, "runner_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest()} + (out / "config.json").write_text(json.dumps(config, indent=2) + "\n") + opener = urllib.request.build_opener(NoRedirect, urllib.request.ProxyHandler({})) + rows = [] + with (out / "results.jsonl").open("w") as stream: + for iteration in range(args.passes): + sequence = [("warmup", entries[i % len(entries)]) for i in range(args.warmup)] + sequence += [("measured", entry) for entry in entries] + for phase, entry in sequence: + row = post(opener, args.base.rstrip("/") + "/v1/systemone", entry["request"]) + row.update(id=entry["id"], phase=phase, pass_index=iteration) + compare(row, references[entry["id"]]) + rows.append(row) + stream.write(json.dumps(row, allow_nan=False) + "\n") + stream.flush() + if phase == "warmup" and not row["ok"]: + break + if phase == "warmup" and not row["ok"]: + break + measured = [row for row in rows if row["phase"] == "measured"] + latencies = [row["client_s"] for row in measured if row["ok"]] + diffs = [row["max_diff"] for row in measured if row["max_diff"] is not None] + summary = {"ok": bool(measured) and all(row["ok"] for row in rows), "concurrency": 1, + "measured": len(measured), "warmup": len(rows) - len(measured), + "failures": sum(not row["ok"] for row in rows), "max_diff": max(diffs, default=None), + "mean_s": statistics.fmean(latencies) if latencies else None, + "median_s": statistics.median(latencies) if latencies else None} + (out / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary)) + return 0 if summary["ok"] else 1 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except (OSError, KeyError, TypeError, ValueError) as error: + sys.exit(f"replay failed: {error}") diff --git a/recipe/jev_vl/requirements.txt b/recipe/jev_vl/requirements.txt new file mode 100644 index 0000000..473a75c --- /dev/null +++ b/recipe/jev_vl/requirements.txt @@ -0,0 +1,9 @@ +# Direct preparation dependencies observed in the historical H800 environment. +# A fresh install of this set has not been validated; this is not a transitive lock. +--extra-index-url https://download.pytorch.org/whl/cu130 +torch==2.13.0+cu130 +transformers==5.17.0 +safetensors==0.8.0 +tokenizers==0.23.2 +huggingface_hub==1.33.0 +pillow==12.3.0 diff --git a/recipe/jev_vl/validation.md b/recipe/jev_vl/validation.md new file mode 100644 index 0000000..c27db15 --- /dev/null +++ b/recipe/jev_vl/validation.md @@ -0,0 +1,103 @@ +# JEV-27B-VL validation + +## Historical H800 validation + +The ABI 6 candidate was measured on one H800 80 GB with CUDA 13.0.88, +driver 580.159.03 and BF16 language weights from +`autotrust/JEV-27B-VL@f34b598d4ef4bcefd337bee8d8e7ddd3b7733ccc`. +The frozen corpus has 36 synthetic text requests and 12 questions about one +synthetic image. It checks implementation parity, not general model quality. + +Jobs 413587 and 413597 tested the same library and worker binaries. The measured +source snapshot predates comment, test-registration and diagnostic-header edits; +its hashes and those of the binaries are in the +[provenance record](https://github.com/linear3735/system1-omni/blob/522f2256876a62ddd70873ee9e775055416aea1b/recipe/jev_vl/evidence/review-20261006/provenance.json). +These are historical measurements, not a new GPU test of later revisions. + +- 12 GPU/kernel and checkpoint checks passed. +- Cache off/on each matched 48/48 reference decisions. Maximum probability + differences were 0.021545 / 0.014215, within the fixed 0.025 tolerance. +- Worker/frontend replay passed 192 measured decisions. Eight worker and seven + supported frontend error probes passed. The frontend owns unsupported-route + 404 responses; the original failed comparison is retained in the archive. + +Each cache mode ran two passes of the same 12 image questions on the same GPU +and binary, at concurrency 1, with 12 excluded warmups before each pass: + +| Mode | Pass 1 p50 | Pass 2 p50 | Combined p50 | +| --- | ---: | ---: | ---: | +| Cache off | 102.35 ms | 104.32 ms | 102.88 ms | +| Cache hit | 37.32 ms | 36.78 ms | 37.04 ms | + +The combined-median ratio is **2.78×**. Timing includes worker-direct localhost +HTTP and response JSON parsing. It excludes startup, offline vision encoding, +warmup and the frontend. Peak memory and multi-client performance were not +measured. Raw timings, cache counters, verdicts and the initial failed comparison +are in the [fixed result archive](https://github.com/linear3735/system1-omni/tree/522f2256876a62ddd70873ee9e775055416aea1b/recipe/jev_vl/evidence/review-20261006). +[Earlier measurements and their limitations](https://github.com/linear3735/system1-omni/blob/4927d2a3373ae3e2d825bf63a320178988540031/recipe/jev_vl/validation.md#appendix-historical-evidence-audited-on-2026-10-06) +remain archived separately. + +## Download the frozen corpus + +Inputs and reference responses are retained at a fixed archive revision rather +than in the current source tree. Download the public archive and extract only +its evidence directory into a temporary directory: + +```sh +jev_vl_evidence=$(mktemp -d "${TMPDIR:-/tmp}/jev-vl-evidence.XXXXXX") +jev_vl_revision=4927d2a3373ae3e2d825bf63a320178988540031 +curl --fail --location \ + "https://codeload.github.com/linear3735/system1-omni/tar.gz/$jev_vl_revision" \ + --output "$jev_vl_evidence/source.tar.gz" +tar -xzf "$jev_vl_evidence/source.tar.gz" -C "$jev_vl_evidence" \ + --strip-components=3 "system1-omni-$jev_vl_revision/recipe/jev_vl/evidence" +python3 "$jev_vl_evidence/evidence/restore_manifest.py" \ + --out "$jev_vl_evidence/manifest.jsonl" +``` + +The restore script verifies the original manifest SHA-256 +`f72d1beaaaf53933d0a6edda26b635f46931990d8cdb56ca5c6a7ca94d2eb0ee`. +It restores approximately 41 MiB of inputs. The extracted `reference/sha256.json` +records hashes of the 48 reference answers and eight archived error probes. +These inputs and responses do not establish the provenance of a new executable. + +## Replay + +Follow the [deployment recipe](README.md) to export weights, build the worker, +preencode `"$jev_vl_evidence/manifest.jsonl"`, and start serving. Keep the shell +variable from the download step. Use a new output directory for every run: + +```sh +python3 recipe/jev_vl/replay.py \ + --base http://127.0.0.1:8001 --manifest "$jev_vl_evidence/manifest.jsonl" \ + --reference "$jev_vl_evidence/evidence/reference" \ + --out "$jev_vl_evidence/worker-parity" --warmup 4 --passes 1 +``` + +The replay tool returns nonzero for failed requests, changed decisions, invalid +probabilities or absolute probability drift above 0.025. Repeat through the +frontend at `http://127.0.0.1:8080` with a different output directory. Archived +error probes are separate from this replay; validate error handling separately. + +For a paired cache measurement, start the worker with `JEV_VL_CACHE=0`, then run: + +```sh +python3 recipe/jev_vl/replay.py \ + --base http://127.0.0.1:8001 --manifest "$jev_vl_evidence/manifest.jsonl" \ + --reference "$jev_vl_evidence/evidence/reference" --pattern img- \ + --out "$jev_vl_evidence/cache-off" --warmup 12 --passes 2 +``` + +Restart the same binary with `JEV_VL_CACHE=1` and repeat into `cache-on`. Keep the +GPU, checkpoint, prepared image assets and request order fixed. Save +`GET /v1/cache/stats` before and after each run to check cache use; counters +include warmup. Keep generated outputs outside the source tree. + +## Coverage limits + +This corpus does not cover general model quality, changed-image workloads, +cache eviction, multi-client performance or full Cua-S1/Open-Jev checkpoint +regressions. Clean installation remains unverified. Use the +[PR checks and review](https://github.com/ThinkFlowLab/system1-omni/pull/96) +for revision-specific CI and review status; passing this corpus does not replace +request-validation tests or a review of later code changes. diff --git a/recipe/open_jev/native.md b/recipe/open_jev/native.md index 245841e..bedf415 100644 --- a/recipe/open_jev/native.md +++ b/recipe/open_jev/native.md @@ -45,7 +45,7 @@ or an incomplete export. The saved limit defaults to 4096 tokens per candidate; The CUDA kernels require compute capability 8.0 or newer. The current build target below is Ada (`89`); pass your GPU's compute capability explicitly. The CUDA shared library and Rust workers must be rebuilt together for ABI -version 5, which includes the shared vision and CUDA Graph entry points. +version 6, which includes vision, CUDA Graph and prefix-continuation entry points. ```sh src/backends/cuda/qwen3_5/build.sh target/release 89 diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index e142176..6382553 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -11,12 +11,19 @@ The norm, elementwise and q/k preparation kernels round to bfloat16 where Transf `cs1_attention_gated` fuses the sigmoid gate into the attention epilogue, preserving the BF16 rounding of both attention and sigmoid before multiplication. The native workers use this entry point; the separate operations remain available for kernel -comparisons. Rebuild the library and workers together for ABI version 5, which -includes the shared vision and CUDA Graph entry points alongside gated attention. +comparisons. Rebuild the library and workers together for ABI version 6, which +includes vision, CUDA Graph, gated attention and prefix-continuation entry points. Gated DeltaNet preparation stores converted TF32 operands in three-byte component planes, preserves the original four-term TF32 accumulation, and writes U/W fragments directly as bfloat16. Dynamic shared memory is 72 KiB per block. The [H200 comparison](../../../../benchmarks/gdn/README.md) records complete GDN call latency, numerical checks, and the small end-to-end change measured with the Open-Jev worker from PR #55. The shared Rust model can pack independent sequences for input and gate/up GEMMs. Output/down GEMMs retain each prompt's original shape and reduction order; -attention, convolution and GDN calls remain sequence-local. The CUDA ABI is -unchanged. Open-Jev uses this path within requests; Cua-S1 keeps single-prompt calls. +attention, convolution and GDN calls remain sequence-local. Packing adds no CUDA +entry points. Open-Jev uses this path within requests; Cua-S1 keeps single-prompt calls. + +Prefix continuation retains full-attention KV, three convolution input rows, +and FP32 Gated DeltaNet state at 64-token chunk boundaries. It is used by the +experimental JEV-VL worker; existing full-prefill callers keep their own paths. +ABI 6 adds mandatory copy and prefix entry points after ABI 5 native vision. +Rebuild the library and **all** Qwen worker binaries together; older libraries +are rejected. Current-branch GPU regression is required before release. diff --git a/src/backends/cuda/qwen3_5/attention.cu b/src/backends/cuda/qwen3_5/attention.cu index 20ffcee..71cf6e8 100644 --- a/src/backends/cuda/qwen3_5/attention.cu +++ b/src/backends/cuda/qwen3_5/attention.cu @@ -78,17 +78,20 @@ constexpr int D = 256, BM = 64, BN = 32, THREADS = 128; constexpr int LDS = D + 8; // shared row stride in elements: 528 bytes keeps ldmatrix conflict-free constexpr int SMEM_BYTES = (BM + 2 * BN) * LDS * 2; +// q_base offsets row 0 of the q/gate/out buffers: this launch covers the window +// [q_base, T_total) while k/v cover [0, T_total). Used by the cached-prefix +// continuation; q_base = 0 reduces to the plain full pass. template __global__ void __launch_bounds__(THREADS) flash_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv, const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk, - float scale_log2) { + float scale_log2, int q_base) { extern __shared__ __align__(16) unsigned char smem[]; bf16* qs = reinterpret_cast(smem); bf16* ks = qs + BM * LDS; bf16* vs = ks + BN * LDS; const int h = blockIdx.y, hk = h / (Hq / Hk); - const int q0 = (gridDim.x - 1 - blockIdx.x) * BM; // the longest blocks first + const int q0 = q_base + (gridDim.x - 1 - blockIdx.x) * BM; // the longest blocks first const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32; const int g = lane / 4, t = lane % 4; const int row0 = q0 + warp * 16; // this warp's first query @@ -232,20 +235,21 @@ __global__ void __launch_bounds__(THREADS) template int launch(const void* q, const void* k, const void* v, int ldv, const void* gate, void* out, - int T, int Hq, int Hk, int Dh, float scale, void* stream) { + int T, int Hq, int Hk, int Dh, float scale, int q_base, void* stream) { if (Dh != D || Hk <= 0 || Hq <= 0 || Hq % Hk != 0 || ldv % 8 != 0 || ldv < Hk * Dh || T < 0) return cudaErrorInvalidValue; - if (T == 0) return cudaSuccess; + if (q_base < 0 || q_base % BM || q_base > T) return cudaErrorInvalidValue; + if (T == q_base) return cudaSuccess; if (Gated && gate == nullptr) return cudaErrorInvalidValue; // Once per specialization (for the device current at the first call). static const cudaError_t configured = cudaFuncSetAttribute( flash_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_BYTES); if (configured != cudaSuccess) return configured; constexpr float LOG2E = 1.4426950408889634f; - flash_kernel<<<<(stream)>>>( static_cast(q), static_cast(k), static_cast(v), ldv, - static_cast(gate), static_cast(out), T, Hq, Hk, scale * LOG2E); + static_cast(gate), static_cast(out), T, Hq, Hk, scale * LOG2E, q_base); return cudaGetLastError(); } @@ -274,10 +278,18 @@ extern "C" int cs1_attn_prep(const void* qg, const void* kr, int ld, const void* extern "C" int cs1_attention(const void* q, const void* k, const void* v, int ldv, void* out, int T, int Hq, int Hk, int Dh, float scale, void* stream) { - return flash::launch(q, k, v, ldv, nullptr, out, T, Hq, Hk, Dh, scale, stream); + return flash::launch(q, k, v, ldv, nullptr, out, T, Hq, Hk, Dh, scale, 0, stream); } extern "C" int cs1_attention_gated(const void* q, const void* k, const void* v, int ldv, const void* gate, void* out, int T, int Hq, int Hk, int Dh, float scale, void* stream) { - return flash::launch(q, k, v, ldv, gate, out, T, Hq, Hk, Dh, scale, stream); + return flash::launch(q, k, v, ldv, gate, out, T, Hq, Hk, Dh, scale, 0, stream); +} + +// Same layout contract as the full pass; only the covered query rows differ: +// q_base is a multiple of 64 and rows below it keep whatever was already in k/v. +extern "C" int cs1_attention_gated_prefix(const void* q, const void* k, const void* v, int ldv, const void* gate, + void* out, int T, int Hq, int Hk, int Dh, float scale, int q_base, + void* stream) { + return flash::launch(q, k, v, ldv, gate, out, T, Hq, Hk, Dh, scale, q_base, stream); } diff --git a/src/backends/cuda/qwen3_5/gdn_prefill.cu b/src/backends/cuda/qwen3_5/gdn_prefill.cu index 676ce44..8355a77 100644 --- a/src/backends/cuda/qwen3_5/gdn_prefill.cu +++ b/src/backends/cuda/qwen3_5/gdn_prefill.cu @@ -337,7 +337,8 @@ constexpr int SS_LD = BVS + 8; // bfloat16 row stride of the S copy and v_n constexpr int STAGE = C * WS_LD; // elements of one staged w or kd constexpr size_t SMEM2_BYTES = (4 * STAGE + K * SS_LD + C * SS_LD) * 2; -__global__ void __launch_bounds__(ST_THREADS) gdn_chunk_state(Work ws, int NC) { +__global__ void __launch_bounds__(ST_THREADS) gdn_chunk_state(Work ws, int NC, const float* __restrict__ S_IN, + float* __restrict__ S_OUT) { extern __shared__ __align__(128) unsigned char sm[]; bf16* wbuf = reinterpret_cast(sm); // [2][C][WS_LD] bf16* kbuf = wbuf + 2 * STAGE; // [2][C][WS_LD] @@ -362,6 +363,21 @@ __global__ void __launch_bounds__(ST_THREADS) gdn_chunk_state(Work ws, int NC) { for (int mt = 0; mt < 2; mt++) #pragma unroll for (int nt = 0; nt < 4; nt++) st[mt][nt][0] = st[mt][nt][1] = st[mt][nt][2] = st[mt][nt][3] = 0.f; + if (S_IN != nullptr) { + // float32 state seeded for a continuation: identical to the registers a + // full pass would carry into this chunk, so the scan repeats one pass. + const float* si = S_IN + (size_t)h * K * V; +#pragma unroll + for (int mt = 0; mt < 2; mt++) +#pragma unroll + for (int nt = 0; nt < 4; nt++) +#pragma unroll + for (int e = 0; e < 4; e++) { + const int row = warp * 32 + mt * 16 + g + (e >> 1) * 8; + const int col = vb0 + nt * 8 + 2 * t + (e & 1); + st[mt][nt][e] = si[(size_t)row * V + col]; + } + } load(0, 0); for (int c = 0; c < NC; c++) { @@ -438,6 +454,20 @@ __global__ void __launch_bounds__(ST_THREADS) gdn_chunk_state(Work ws, int NC) { } } } + if (S_OUT != nullptr) { + float* so = S_OUT + (size_t)h * K * V; +#pragma unroll + for (int mt = 0; mt < 2; mt++) +#pragma unroll + for (int nt = 0; nt < 4; nt++) +#pragma unroll + for (int r = 0; r < 2; r++) { + const int row = warp * 32 + mt * 16 + g + r * 8; + const int col = vb0 + nt * 8 + 2 * t; + *reinterpret_cast(so + (size_t)row * V + col) = + make_float2(st[mt][nt][2 * r], st[mt][nt][2 * r + 1]); + } + } } // ---- kernel 3: the output, per chunk ---- @@ -535,10 +565,20 @@ extern "C" { size_t cs1_gdn_workspace_floats(int T, int H) { return (Layout(T, H).total + 3) / 4; } -int cs1_gdn_prefill(const void* q, const void* k, const void* v, const float* g, const void* beta, - void* o, float* workspace, int T, int H, int HK, float scale, void* stream) { +static int gdn_prefill_run(const void* q, const void* k, const void* v, const float* g, + const void* beta, void* o, float* workspace, int T, int H, int HK, + float scale, const void* s_in, void* s_out, void* stream) { if (T < 0 || HK <= 0 || H % HK != 0) return cudaErrorInvalidValue; - if (T == 0) return cudaSuccess; + const float* si = static_cast(s_in); + float* so = static_cast(s_out); + if (T == 0) { + // No tokens: the state passes through unchanged, for symmetric capture. + if (so == nullptr) return cudaSuccess; + const size_t bytes = (size_t)H * K * V * sizeof(float); + cudaStream_t st = static_cast(stream); + return si != nullptr ? cudaMemcpyAsync(so, si, bytes, cudaMemcpyDeviceToDevice, st) + : cudaMemsetAsync(so, 0, bytes, st); + } // once per process (for the device current at the first call) static const cudaError_t configured = [] { cudaError_t e = cudaFuncSetAttribute(gdn_chunk_prep, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -558,9 +598,20 @@ int cs1_gdn_prefill(const void* q, const void* k, const void* v, const float* g, gdn_chunk_prep<<>>( static_cast(q), static_cast(k), static_cast(v), g, static_cast(beta), ws, T, H, HK, scale); - gdn_chunk_state<<>>(ws, NC); + gdn_chunk_state<<>>(ws, NC, si, so); gdn_chunk_out<<>>(ws, static_cast(o), T, H); return cudaGetLastError(); } +int cs1_gdn_prefill(const void* q, const void* k, const void* v, const float* g, const void* beta, + void* o, float* workspace, int T, int H, int HK, float scale, void* stream) { + return gdn_prefill_run(q, k, v, g, beta, o, workspace, T, H, HK, scale, nullptr, nullptr, stream); +} + +int cs1_gdn_prefill_x(const void* q, const void* k, const void* v, const float* g, const void* beta, + void* o, float* workspace, int T, int H, int HK, float scale, + const void* s_in, void* s_out, void* stream) { + return gdn_prefill_run(q, k, v, g, beta, o, workspace, T, H, HK, scale, s_in, s_out, stream); +} + } // extern "C" diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 37a08a7..50d3618 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -13,7 +13,7 @@ #include // Bumped whenever the required interface below changes. -#define CS1_ABI_VERSION 5 +#define CS1_ABI_VERSION 6 #ifdef __cplusplus extern "C" { @@ -36,6 +36,11 @@ int cs1_graph_destroy(void* exec); // Copy and wait for the copy. int cs1_upload(void* dst, const void* src, size_t bytes, void* stream); int cs1_download(void* dst, const void* src, size_t bytes, void* stream); +// Device-to-device copy of `bytes` (queue-only completion). +int cs1_copy_dd(void* dst, const void* src, size_t bytes, void* stream); +// Pitched copy both ways (default direction = device-to-device): +// height rows of `width`, output pitch dpitch, input pitch spitch, in bytes. +int cs1_copy2d(void* dst, size_t dpitch, const void* src, size_t spitch, size_t width, size_t height, void* stream); // ---- operations ---- @@ -70,6 +75,15 @@ size_t cs1_gdn_workspace_floats(int T, int H); int cs1_gdn_prefill(const void* q, const void* k, const void* v, const float* g, const void* beta, void* o, float* workspace, int T, int H, int HK, float scale, void* stream); +// Same scan with an explicit state: s_in is the [H, K, V] float32 state before +// this window (null = zero), s_out receives the [H, K, V] float32 state after +// it (null = skip). The float32 state makes a continuation starting exactly at +// a chunk boundary repeat the arithmetic of one full pass bit for bit. Used +// for the cached-prefix continuation and for its capture. +int cs1_gdn_prefill_x(const void* q, const void* k, const void* v, const float* g, const void* beta, + void* o, float* workspace, int T, int H, int HK, float scale, + const void* s_in, void* s_out, void* stream); + // Attention inputs: q and gate from qg [T, Hq, 2*Dh], k from kr [T, Hk, Dh], both in rows // of ld; per-head zero-centred RMSNorm, then rotary embedding on the first 2*half dims // using cos/sin [T, half] (bfloat16). Writes q [T, Hq, Dh], gate [T, Hq*Dh], k [T, Hk, Dh]. @@ -88,6 +102,13 @@ int cs1_attention(const void* q, const void* k, const void* v, int ldv, void* ou int cs1_attention_gated(const void* q, const void* k, const void* v, int ldv, const void* gate, void* out, int T, int Hq, int Hk, int Dh, float scale, void* stream); +// Windowed variant for the cached prefix: k and v cover rows [0, T) (the cached +// prefix must already be in place); q, gate and out cover rows [q_base, T) only. +// q_base must be a multiple of 64; q_base = 0 reduces to cs1_attention_gated. +int cs1_attention_gated_prefix(const void* q, const void* k, const void* v, int ldv, const void* gate, + void* out, int T, int Hq, int Hk, int Dh, float scale, int q_base, + void* stream); + // x = x * sigmoid(gate), n elements. int cs1_sigmoid_gate(void* x, const void* gate, size_t n, void* stream); diff --git a/src/backends/cuda/qwen3_5/runtime.cu b/src/backends/cuda/qwen3_5/runtime.cu index 697c9db..5ac61ce 100644 --- a/src/backends/cuda/qwen3_5/runtime.cu +++ b/src/backends/cuda/qwen3_5/runtime.cu @@ -36,6 +36,19 @@ int cs1_download(void* dst, const void* src, size_t bytes, void* stream) { return e != cudaSuccess ? e : cudaStreamSynchronize(st); } +int cs1_copy_dd(void* dst, const void* src, size_t bytes, void* stream) { + const cudaError_t e = + cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, static_cast(stream)); + return e != cudaSuccess ? e : cudaGetLastError(); +} + +int cs1_copy2d(void* dst, size_t dpitch, const void* src, size_t spitch, size_t width, size_t height, + void* stream) { + const cudaError_t e = cudaMemcpy2DAsync(dst, dpitch, src, spitch, width, height, + cudaMemcpyDeviceToDevice, static_cast(stream)); + return e != cudaSuccess ? e : cudaGetLastError(); +} + int cs1_graph_begin(void* stream) { const cudaError_t e = cudaStreamBeginCapture(static_cast(stream), cudaStreamCaptureModeThreadLocal); if (e != cudaSuccess) (void)cudaGetLastError(); diff --git a/src/models/jev_vl/README.md b/src/models/jev_vl/README.md new file mode 100644 index 0000000..c6b3c70 --- /dev/null +++ b/src/models/jev_vl/README.md @@ -0,0 +1,104 @@ +# JEV-27B-VL experimental native worker + +This integration targets **autotrust/JEV-27B-VL**, with a native Rust/CUDA +language backbone and an offline Hugging Face vision encoder. The reviewed +integrated candidate passed the frozen-corpus H800 checks; it remains experimental +and has not completed broad quality or clean-install validation. See the [recipe](../../../recipe/jev_vl/README.md) +and [evidence and limitations](../../../recipe/jev_vl/validation.md). + +## Pinned reference and ownership + +The checkpoint, tokenizer, `adapter_vllm`, calibration and `serve_decide.py` +reference are from +[autotrust/JEV-27B-VL at f34b598d4ef4bcefd337bee8d8e7ddd3b7733ccc](https://huggingface.co/autotrust/JEV-27B-VL/tree/f34b598d4ef4bcefd337bee8d8e7ddd3b7733ccc). +The model card declares Apache-2.0. Prompt rendering and verbalizer selection +follow that release's `serve_decide.py`; source attribution, adapted portions and +the license are recorded in the [third-party notices](native/THIRD_PARTY_NOTICES.md). +The exporter records checkpoint-index, adapter, decision-head and calibration +hashes in `jev_vl_export.json`. + +The language executor reuses the shared Qwen3.5/3.8 implementation established by +[Open-Jev PR #55](https://github.com/ThinkFlowLab/system1-omni/pull/55). +This checkpoint's verbalizer head and prompt differ from +ZefanCai/Open-Jev-27B-v1.1's independent-candidate scalar head. It is also distinct +from [openjev/openjev issue #95](https://github.com/ThinkFlowLab/system1-omni/issues/95). +[Cua-S1 PR #64](https://github.com/ThinkFlowLab/system1-omni/pull/64) supplies +upstream native vision infrastructure; this JEV worker currently consumes +preencoded embeddings and does not call that native vision tower. +[Shared-observation RFC #85](https://github.com/ThinkFlowLab/system1-omni/issues/85) +is related design context, not this model's acceptance specification. + +## Request and response contract + +`POST /v1/systemone` accepts the reference's **single-question** shape: + +```json +{"kind":"choice","state":"The package arrived damaged.","question":"What should support do?","options":["Offer a replacement","Close the ticket"]} +``` + +This differs from the `questions`/`answers` envelope used by other workers. The +frontend transports this body unchanged; generic clients must use this worker's +schema. It exposes only the health and decision routes; unsupported generation +routes receive the frontend's own 404. An optional `model` must equal +`autotrust/JEV-27B-VL`. + +| Kind | Candidate set and output | +| --- | --- | +| `choice` | 2–256 string options in request order; the first 16 use trained A–P slots, later labels use the exported single-token table with zero extra slot bias. | +| `noul` | Fixed options `"false"`, `"true"`. | +| `score` | Fixed options `"0"` through `"5"`; the selected score is the highest-probability category, not an expected value. | + +Responses include `kind`, `effective_kind`, `options`, `probabilities`, +`choice_index`, `choice`, `adaptation`, `protocol`, `model`, `usage`, +`elapsed_seconds` and `num_model_requests`. Ties select the earliest option. +`usage` counts the full expanded prompt even on cache hits and reports one +completion token for the decision; no autoregressive decoding takes place. + +The raw prompt is `[kind] …\n[state] …\n[question] …\n[options]\n…\n[decision]:`. +It uses no chat template. Structured state preserves Python-style JSON rendering; +list parts concatenate without an inserted separator. Each image part inserts +`<|vision_start|><|image_pad|><|vision_end|>` before image-token expansion. + +The default exported limit is 16,384 expanded tokens. Export permits 1–32,768; +larger limits have not been benchmarked. HTTP bodies are limited to 4 MiB. +`thinking=auto/on`, tournament and permutation strategies, generation, audio, +video, dynamic batching and unprepared images are unsupported. `strategy=auto` +uses the single-pass path in this worker. The historical corpus covers only +2–16 choice options and one image, so the wider accepted range and multiple +images do not have full-checkpoint parity evidence. + +## Execution, layouts and lifetimes + +`processing.rs` validates and tokenizes one question and prepares either text +IDs, full multimodal inputs, or a cached-prefix continuation. Image assets hold +BF16 `[image_tokens, 5120]` adapted embeddings and their patch grid. The +processor constructs expanded token IDs and three position axes; image rows +replace placeholder-token embeddings. All vectors describe one unpadded prompt, +not a GPU batch. + +`executor.rs` owns checkpoint loading, the model forward and selected output-head +rows. The final 5120-element hidden state is downloaded and projected against +FP32 merged head rows using FP64 accumulation. Per-kind probabilities are +`softmax((label_logit + slot_bias) / temperature)`; the common vocabulary +log-normalizer cancels. Response assembly stays in the processor. This is a +mathematical readout equivalence, not a claim of bitwise equivalence to vLLM. + +The shared `SerialScheduler` admits one complete question per executor. A model +mutex protects mutable CUDA state; the scheduler retains admission until the +blocking forward and readout finish. Model-owned prefix snapshots contain +full-attention K/V, FP32 Gated DeltaNet state and convolution history. A +continuation owns references to its snapshot and image assets until execution +finishes. Restart the worker and use a fresh asset directory when changing +weights, tokenizer or preprocessing; caches are scoped to a loaded worker. + +Three configurable caches retain processor prefix geometry (L1), parsed +preencoded image assets (L2), and language-prefix device state (L3). The current +L3 path targets a compatible single-image prefix aligned to a 64-token chunk; +other inputs use a full forward. L2 avoids file parsing and copying of prepared +assets; it does not run or cache a live vision encoder. Cache sizes and controls +are listed in the [recipe](../../../recipe/jev_vl/README.md#cache-controls). + +Tests are registered under [`tests/jev_vl/`](../../../tests/jev_vl/) and +[`tests/qwen3_5/`](../../../tests/qwen3_5/). CPU tests establish contract and +bookkeeping behavior; ignored checkpoint/kernel tests and paired full-model +validation are separate gates. diff --git a/src/models/jev_vl/native/Cargo.toml b/src/models/jev_vl/native/Cargo.toml new file mode 100644 index 0000000..84bd8c8 --- /dev/null +++ b/src/models/jev_vl/native/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "omni-jev-vl-native" +version = "0.1.0" +edition = "2024" +publish = false +description = "Native Rust/CUDA System-1 worker for autotrust/JEV-27B-VL (verbalizer readout)" + +[dependencies] +anyhow = "1.0.100" +axum = "0.8.8" +half = "2.7.1" +sha2 = "0.11" +omni-qwen3-5-native = { path = "../../qwen3_5/native" } +omni-runtime = { path = "../../../runtime" } +safetensors = "0.8.0" +serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order"] } +tokenizers = { version = "=0.22.2", default-features = false, features = ["onig"] } +tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "sync"] } + +[[test]] +name = "jev_vl_contract" +path = "../../../../tests/jev_vl/contract.rs" + +[[test]] +name = "jev_vl_images" +path = "../../../../tests/jev_vl/images.rs" + +[[test]] +name = "jev_vl_prefix" +path = "../../../../tests/jev_vl/prefix.rs" + +[dev-dependencies] +tempfile = "3" diff --git a/src/models/jev_vl/native/THIRD_PARTY_NOTICES.md b/src/models/jev_vl/native/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..6a66a8e --- /dev/null +++ b/src/models/jev_vl/native/THIRD_PARTY_NOTICES.md @@ -0,0 +1,30 @@ +# JEV-27B-VL adaptation notices + +This worker and its preparation scripts use the following pinned references. +The applicable license text is retained as +[Apache License, Version 2.0](licenses/APACHE-2.0). + +- **autotrust/JEV-27B-VL**, revision + `f34b598d4ef4bcefd337bee8d8e7ddd3b7733ccc`: + [`serve_decide.py`](https://huggingface.co/autotrust/JEV-27B-VL/blob/f34b598d4ef4bcefd337bee8d8e7ddd3b7733ccc/serve_decide.py). + `src/contract.rs` adapts its state rendering, raw System-1 prompt and decision + response semantics to Rust. `recipe/jev_vl/export_merged.py` reproduces its + single-token option-label scan. The Rust worker supports a restricted + single-question, System-1-only interface and does not include the upstream + server, adaptive thinking or multi-pass strategies. The pinned + [model card](https://huggingface.co/autotrust/JEV-27B-VL/blob/f34b598d4ef4bcefd337bee8d8e7ddd3b7733ccc/README.md) + declares Apache-2.0 for the release; the included license text comes from its + `LICENSE` file. No weights are redistributed with this integration. +- **Transformers 5.17.0**: + [`modeling_qwen3_5.py`](https://github.com/huggingface/transformers/blob/v5.17.0/src/transformers/models/qwen3_5/modeling_qwen3_5.py), + Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. + `src/images.rs` adapts the multimodal position-index calculation to Rust, + preencoded image rows and bounded grid dimensions. + [`vision_utils.py`](https://github.com/huggingface/transformers/blob/v5.17.0/src/transformers/vision_utils.py), + Copyright 2026 The HuggingFace Inc. team. All rights reserved, is the related + vision-position reference. These sources are licensed under Apache-2.0. + +The offline preencoder invokes Transformers, PyTorch and Pillow as installed +dependencies. It does not vendor their vision model or image decoder code. +Existing shared Cua-S1 native vision adaptations retain their separate +[notices](../../cua_s1/native/THIRD_PARTY_NOTICES.md). diff --git a/src/models/jev_vl/native/licenses/APACHE-2.0 b/src/models/jev_vl/native/licenses/APACHE-2.0 new file mode 100644 index 0000000..1d5180a --- /dev/null +++ b/src/models/jev_vl/native/licenses/APACHE-2.0 @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2026 Alibaba Cloud + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/src/models/jev_vl/native/src/caches.rs b/src/models/jev_vl/native/src/caches.rs new file mode 100644 index 0000000..99b0b4a --- /dev/null +++ b/src/models/jev_vl/native/src/caches.rs @@ -0,0 +1,384 @@ +//! Process-local caches for prepared images and reusable language prefixes. +//! +//! - L1 processor: a per-prefix-structure record. Key = sha256 over the exact +//! request structure (kind + every state part: text verbatim / image url). +//! Value = the compiled prefix geometry (pad run, cached-prefix length P, +//! position bases, the expanded prefix ids/positions), so a hit skips the +//! prompt-head assembly, tokenization, expansion and meshgrid work; the suffix +//! (`<|vision_end|>` onward) is always tokenized fresh per question. +//! - L2 vision: parsed image assets (adapter output rows + grid) keyed by +//! sha256(url) — the same key form as the offline imgcache directory. A hit +//! skips the disk read, safetensors parsing and validation. Vision encoding +//! runs offline and is not part of this cache lookup. +//! - L3 KV prefix: the device-side prefix state (per-full-attention-layer +//! post-prep K and raw V rows; per-GDN-layer float32 recurrent state + conv +//! tail). Holders are `Arc`; a continuation reads it while the +//! device buffers stay alive. Full-attention KV is per-prefix; the GDN/conv +//! states are per-prefix too — identical token prefix ⇒ identical state, and +//! they never blend across different prefixes (hybrid constraint). +//! +//! All collections are process-local, LRU + budgeted; everything is disabled by +//! JEV_VL_CACHE=0 and per level by JEV_VL_L1/L2/L3=0 (A/B decomposition). + +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex, OnceLock}; + +use omni_qwen3_5_native::model::PrefixState; + +use crate::images::ImageAsset; + +/// Env-read switches + budgets. +pub struct CacheCfg { + pub enabled: bool, + pub l1: bool, + pub l2: bool, + pub l3: bool, + /// Max structural records. + pub l1_max: usize, + /// L2 parsed-asset budget in bytes. + pub l2_bytes: usize, + /// L3 retained device-state budget in bytes; zero disables retention. + pub l3_bytes: usize, +} + +impl CacheCfg { + pub fn from_env() -> Self { + let flag = |name: &str, default: bool| { + std::env::var(name).map_or(default, |v| v != "0" && !v.eq_ignore_ascii_case("false")) + }; + let num = |name: &str, default: usize| { + std::env::var(name) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) + }; + Self { + enabled: flag("JEV_VL_CACHE", true), + l1: flag("JEV_VL_L1", true), + l2: flag("JEV_VL_L2", true), + l3: flag("JEV_VL_L3", true), + l1_max: num("JEV_VL_L1_MAX", 256), + l2_bytes: num("JEV_VL_L2_BYTES", 1 << 30), + l3_bytes: num("JEV_VL_L3_BYTES", 2 << 30), + } + } +} + +#[derive(Clone, Default)] +pub struct CacheStats { + pub l1_hit: u64, + pub l1_miss: u64, + pub l2_hit: u64, + pub l2_miss: u64, + pub l3_hit: u64, + pub l3_miss: u64, + pub l3_populate: u64, + pub l3_fallback_full: u64, + pub l1_records: u64, + pub l1_bytes: u64, + pub l2_records: u64, + pub l2_bytes: u64, + pub l3_records: u64, + pub l3_bytes: u64, +} + +/// The compiled geometry of one structural prefix (small; see module doc). +pub struct L1Meta { + /// Global row of the last image block's first pad; also its position base. + pub pads_start: usize, + /// Global row one past the last image block's last pad. + pub pads_end: usize, + /// Cached-prefix length, a multiple of 64: floor(pads_end / 64) * 64. + pub p: usize, + /// Position base of the pads (text-position base for the meshgrid offset). + pub base_pad: i64, + /// mrope advance after the image: max(grid_h, grid_w) / 2. + pub advance: i64, + /// Expanded ids of rows [0, p) (needed by the L1-only mode's full rebuild). + pub ids_prefix: Vec, + /// Absolute positions of rows [0, p). + pub positions_prefix: [Vec; 3], +} + +impl L1Meta { + fn bytes(&self) -> usize { + self.ids_prefix.len() * 4 + 3 * self.positions_prefix[0].len() * 8 + 96 + } +} + +struct Registry { + map: HashMap>, + lru: VecDeque, + meta_bytes: usize, + state_bytes: usize, + max_records: usize, + state_budget: usize, +} + +/// One structural record: compiled prefix geometry + (lazily) its device state. +pub struct PrefixRecord { + pub key: u64, + pub meta: L1Meta, + state: OnceLock>, +} + +impl PrefixRecord { + pub fn state(&self) -> Option> { + self.state.get().cloned() + } +} + +/// Parsed image assets (L2). +struct L2Store { + map: HashMap>, + lru: VecDeque, + bytes: usize, + budget: usize, +} + +/// Cache state shared by the processor and executor. +pub struct Caches { + pub cfg: CacheCfg, + l2: Mutex, + registry: Mutex, + stats: Mutex, +} + +impl Caches { + pub fn new(cfg: CacheCfg) -> Arc { + Arc::new(Self { + l2: Mutex::new(L2Store { + map: HashMap::new(), + lru: VecDeque::new(), + bytes: 0, + budget: cfg.l2_bytes, + }), + registry: Mutex::new(Registry { + map: HashMap::new(), + lru: VecDeque::new(), + meta_bytes: 0, + state_bytes: 0, + max_records: cfg.l1_max, + state_budget: cfg.l3_bytes, + }), + cfg, + stats: Mutex::new(CacheStats::default()), + }) + } + + fn bump(&self, f: impl FnOnce(&mut CacheStats)) { + if let Ok(mut s) = self.stats.lock() { + f(&mut s) + } + } + + pub fn snapshot(&self) -> CacheStats { + let mut s = self.stats.lock().map(|s| s.clone()).unwrap_or_default(); + if let Ok(r) = self.registry.lock() { + s.l1_records = r.map.len() as u64; + s.l1_bytes = r.meta_bytes as u64; + s.l3_records = r.map.values().filter(|rec| rec.state().is_some()).count() as u64; + s.l3_bytes = r.state_bytes as u64; + } + if let Ok(l2) = self.l2.lock() { + s.l2_records = l2.map.len() as u64; + s.l2_bytes = l2.bytes as u64; + } + s + } + + pub fn reset_stats(&self) { + if let Ok(mut s) = self.stats.lock() { + *s = CacheStats::default(); + } + } + + // ---- L2: parsed image assets ---- + + pub fn l2_get(&self, key: &str) -> Option> { + if !(self.cfg.enabled && self.cfg.l2) { + self.bump(|s| s.l2_miss += 1); + return None; + } + let hit = self.l2.lock().ok().and_then(|mut st| { + let asset = st.map.get(key).cloned(); + if asset.is_some() { + st.lru.retain(|k| k != key); + st.lru.push_back(key.to_owned()); + } + asset + }); + match &hit { + Some(_) => self.bump(|s| s.l2_hit += 1), + None => self.bump(|s| s.l2_miss += 1), + } + hit + } + + /// Insert a freshly loaded asset, evicting LRU entries over the byte budget. + pub fn l2_insert(&self, key: String, asset: Arc) { + if !(self.cfg.enabled && self.cfg.l2) { + return; + } + let bytes = asset.embeddings.len() * 2 + 64; + if bytes <= self.cfg.l2_bytes + && let Ok(mut st) = self.l2.lock() + { + st.lru.retain(|k| k != &key); + st.lru.push_back(key.clone()); + if let Some(old) = st.map.insert(key, asset) { + st.bytes = st.bytes.saturating_sub(old.embeddings.len() * 2 + 64); + } + st.bytes += bytes; + while st.bytes > st.budget && st.lru.len() > 1 { + let evict = st.lru.pop_front().unwrap(); + if let Some(old) = st.map.remove(&evict) { + st.bytes = st.bytes.saturating_sub(old.embeddings.len() * 2 + 64); + } + } + } + } + + // ---- L1 + L3: structural records ---- + + pub fn record_get(&self, key: u64) -> Option> { + if !(self.cfg.enabled && self.cfg.l1) { + self.bump(|s| s.l1_miss += 1); + return None; + } + let hit = self.registry.lock().ok().and_then(|mut r| { + let rec = r.map.get(&key).cloned(); + if rec.is_some() { + r.lru.retain(|k| *k != key); + r.lru.push_back(key); + } + rec + }); + match &hit { + Some(_) => self.bump(|s| s.l1_hit += 1), + None => self.bump(|s| s.l1_miss += 1), + } + hit + } + + /// Insert a new structural record, evicting LRU records over the record + /// count (evicted device states drop with their Arc). + pub fn record_insert(&self, key: u64, meta: L1Meta) -> Arc { + let rec = Arc::new(PrefixRecord { + key, + state: OnceLock::new(), + meta, + }); + if !(self.cfg.enabled && self.cfg.l1) || self.cfg.l1_max == 0 { + return rec; + } + if let Ok(mut r) = self.registry.lock() { + r.meta_bytes = r.meta_bytes.saturating_add(rec.meta.bytes()); + r.lru.retain(|k| *k != key); + r.lru.push_back(key); + if let Some(old) = r.map.insert(key, rec.clone()) { + r.meta_bytes = r.meta_bytes.saturating_sub(old.meta.bytes()); + if let Some(s) = old.state() { + r.state_bytes = r.state_bytes.saturating_sub(s.bytes()); + } + } + loop { + let over = r.map.len() > r.max_records; + if !over || r.lru.len() <= 1 { + break; + } + let evict = r.lru.pop_front().unwrap(); + if evict == key { + r.lru.push_front(evict); + break; + } + if let Some(old) = r.map.remove(&evict) { + r.meta_bytes = r.meta_bytes.saturating_sub(old.meta.bytes()); + if let Some(s) = old.state() { + r.state_bytes = r.state_bytes.saturating_sub(s.bytes()); + } + } + } + } + rec + } + + /// Publish the device state for a record the executor just populated; + /// Enforce the retained-state budget before publication. No-op when the + /// record was already populated, evicted, or larger than the budget. + pub fn record_publish_state(&self, key: u64, state: PrefixState) { + if !(self.cfg.enabled && self.cfg.l3) { + return; + } + let Ok(mut r) = self.registry.lock() else { + return; + }; + let Some(rec) = r.map.get(&key).cloned() else { + return; + }; + let bytes = state.bytes(); + if rec.state.get().is_some() || bytes > r.state_budget || r.state_budget == 0 { + return; + } + // Publication and eviction share the registry lock. A state is immutable + // once published, so readers never need a second, oppositely ordered lock. + r.lru.retain(|k| *k != key); + r.lru.push_back(key); + while r.state_bytes > r.state_budget - bytes { + let Some(evict) = r.lru.pop_front() else { + return; + }; + if let Some(old) = r.map.remove(&evict) { + r.meta_bytes = r.meta_bytes.saturating_sub(old.meta.bytes()); + if let Some(s) = old.state() { + r.state_bytes = r.state_bytes.saturating_sub(s.bytes()); + } + } + } + let _ = rec.state.set(Arc::new(state)); + r.state_bytes += bytes; + } + + pub fn l3_hit(&self) { + self.bump(|s| s.l3_hit += 1); + } + + pub fn l3_miss(&self) { + self.bump(|s| s.l3_miss += 1); + } + + pub fn l3_populate(&self) { + self.bump(|s| s.l3_populate += 1); + } + + pub fn l3_fallback_full(&self) { + self.bump(|s| s.l3_fallback_full += 1); + } +} + +/// Key of one structural prefix: kind, then every state part verbatim +/// (text body / image url). Two requests share L1..L3 state only through an +/// identical key — the "same image, many questions" anchor. +pub fn structure_key(kind: &str, parts: &[crate::contract::Part]) -> u64 { + use sha2::Digest; + let mut h = sha2::Sha256::new(); + h.update(b"jev27-prefix-v1\0"); + h.update((kind.len() as u32).to_le_bytes()); + h.update(kind.as_bytes()); + for part in parts { + match part { + crate::contract::Part::Text(t) => { + h.update([1u8]); + h.update((t.len() as u64).to_le_bytes()); + h.update(t.as_bytes()); + } + crate::contract::Part::Image(url) => { + h.update([2u8]); + h.update((url.len() as u64).to_le_bytes()); + h.update(url.as_bytes()); + } + } + } + let digest = h.finalize(); + u64::from_le_bytes(digest[..8].try_into().unwrap()) +} diff --git a/src/models/jev_vl/native/src/contract.rs b/src/models/jev_vl/native/src/contract.rs new file mode 100644 index 0000000..f373213 --- /dev/null +++ b/src/models/jev_vl/native/src/contract.rs @@ -0,0 +1,404 @@ +//! JEV-27B-VL System-1 request contract and prompt compilation. +//! +//! The request is the official /v1/decide shape ({kind, state, question, +//! options}); the raw System-1 prompt mirrors serve_decide.py's `s1_pass`: +//! `[kind] {kind}\n[state] {state}\n[question] {question}\n[options]\n{lines}\n[decision]:` +//! with Python `json.dumps(..., ensure_ascii=False)` rendering for structured +//! state parts. Rejections carry the official response envelope +//! ({"error": {message, type, param, code}}) with the upstream status codes, so +//! the frozen boundary probes compare equal under digit normalization. + +use omni_qwen3_5_native::json; +use serde_json::{Map, Value, json}; + +pub const MODEL_ID: &str = "autotrust/JEV-27B-VL"; +pub const PROTOCOL: &str = "jev27-bare-v1"; +pub const IMAGE_PLACEHOLDER: &str = "<|vision_start|><|image_pad|><|vision_end|>"; + +/// A rejection keeping the official status code and body envelope. +#[derive(Debug)] +pub struct Reject { + pub status: u16, + pub body: Value, +} + +impl Reject { + fn new(status: u16, message: String, type_: &str, param: Option<&str>) -> Self { + let mut error = json!({"message": message, "type": type_}); + if let Some(param) = param { + error["param"] = json!(param); + } + error["code"] = json!(status); + Self { + status, + body: json!({"error": error}), + } + } + + /// serve_decide.py's `_err`: semantic rejections from the decide layer. + pub fn bad_request(message: impl Into) -> Self { + Self::new(400, message.into(), "BadRequestError", None) + } + + /// The upstream pydantic-validation envelope ("Bad Request"). + pub fn validation(message: String, param: &str) -> Self { + Self::new(400, message, "Bad Request", Some(param)) + } + + fn missing(field: &str, items: usize) -> Self { + Self::validation( + format!( + "1 validation error:\n {{'type': 'missing', 'loc': 'body.{field}', 'msg': 'Field required', 'input': ''}}" + ), + field, + ) + } + + pub fn unknown_model(model: &str) -> Self { + Self::new( + 404, + format!("The model `{model}` does not exist."), + "NotFoundError", + Some("model"), + ) + } +} + +#[derive(Clone, Copy, PartialEq, Debug)] +pub enum Kind { + Noul, + Score, + Choice, +} + +impl Kind { + pub fn as_str(self) -> &'static str { + match self { + Kind::Noul => "noul", + Kind::Score => "score", + Kind::Choice => "choice", + } + } +} + +#[derive(Clone, Debug)] +pub enum Part { + Text(String), + Image(String), +} + +#[derive(Debug)] +pub struct Compiled { + pub kind: Kind, + pub parts: Vec, + pub question: String, + pub options: Vec, + /// The raw System-1 prompt; image parts become the placeholder text. + pub prompt: String, + pub images: Vec, +} + +/// Python repr() of a JSON scalar, for pydantic-style error interpolation. +fn py_repr(value: &Value) -> String { + match value { + Value::String(s) => { + if !s.contains('\'') { + format!("'{s}'") + } else { + format!("\"{}\"", s.replace('"', "\\\"")) + } + } + Value::Null => "None".into(), + other => json::dumps(other), + } +} + +/// `json.dumps(value, ensure_ascii=False)` for dict/other state parts. +fn render(value: &Value) -> String { + if let Some(text) = value.as_str() { + return text.to_owned(); + } + json::dumps(value) +} + +/// serve_decide.py's `_parts`: state -> text/image parts + has_image. +pub fn parts(state: &Value) -> (Vec, Vec) { + match state { + Value::String(text) => (vec![Part::Text(text.clone())], Vec::new()), + Value::Object(_) => (vec![Part::Text(render(state))], Vec::new()), + Value::Array(list) => { + let mut out = Vec::with_capacity(list.len()); + let mut images = Vec::new(); + for p in list { + match p { + Value::String(text) => out.push(Part::Text(text.clone())), + Value::Object(map) if map.contains_key("image") => { + let url = map["image"].as_str().unwrap_or_default().to_owned(); + images.push(url); + out.push(Part::Image(images.last().unwrap().clone())); + } + Value::Object(map) + if map.get("type").and_then(Value::as_str) == Some("image_url") => + { + let url = map + .get("image_url") + .and_then(|v| v.get("url")) + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(); + images.push(url.clone()); + out.push(Part::Image(url)); + } + Value::Object(map) + if map.get("type").and_then(Value::as_str) == Some("text") => + { + out.push(Part::Text( + map.get("text") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + )); + } + other => out.push(Part::Text(render(other))), + } + } + (out, images) + } + _ => (vec![Part::Text(render(state))], Vec::new()), + } +} + +/// Validate and compile a request; mirrors the /v1/decide entrance checks. +/// `labels` is the exported single-token option-label list (its length is the +/// advertised option bound, 256). +pub fn compile(raw: &[u8], labels: &[String]) -> Result { + let request: Map = match json::parse(raw) { + Ok(map) => map, + Err(_) => { + return Err(Reject::validation( + format!( + "1 validation error:\n {{'type': 'model_attributes_type', 'loc': 'body', 'msg': 'Input should be a valid dictionary or object to extract fields from', 'input': ''}}", + raw.len() + ), + "body", + )); + } + }; + let items = request.len(); + if let Some(model) = request.get("model").and_then(Value::as_str) + && model != MODEL_ID + { + return Err(Reject::unknown_model(model)); + } + let kind = match request.get("kind") { + None => return Err(Reject::missing("kind", items)), + Some(v) => match v.as_str() { + Some("noul") => Kind::Noul, + Some("score") => Kind::Score, + Some("choice") => Kind::Choice, + Some(_) => { + return Err(Reject::validation( + format!( + "1 validation error:\n {{'type': 'literal_error', 'loc': 'body.kind', 'msg': \"Input should be 'noul', 'score' or 'choice'\", 'input': {}, 'ctx': {{'expected': \"'noul', 'score' or 'choice'\"}}}}", + py_repr(v) + ), + "kind", + )); + } + None => { + return Err(Reject::validation( + format!( + "1 validation error:\n {{'type': 'literal_error', 'loc': 'body.kind', 'msg': \"Input should be 'noul', 'score' or 'choice'\", 'input': {}, 'ctx': {{'expected': \"'noul', 'score' or 'choice'\"}}}}", + py_repr(v) + ), + "kind", + )); + } + }, + }; + let question = match request.get("question") { + None => return Err(Reject::missing("question", items)), + Some(Value::String(q)) => q.clone(), + _ => { + return Err(Reject::bad_request("question must be a string".to_owned())); + } + }; + if let Some(thinking) = request.get("thinking") { + let thinking = thinking + .as_str() + .ok_or_else(|| Reject::bad_request("thinking must be a string"))?; + match thinking { + "default" | "off" => {} + "auto" | "on" => { + if kind == Kind::Score { + return Err(Reject::bad_request( + "thinking is supported for noul and choice", + )); + } + return Err(Reject::bad_request( + "adaptive thinking is not implemented by this worker", + )); + } + _ => { + return Err(Reject::bad_request(format!( + "thinking must be one of default, off, auto, on; got {thinking:?}" + ))); + } + } + } + if request + .get("system2_only") + .is_some_and(|value| value != &Value::Bool(false)) + { + return Err(Reject::bad_request( + "system2_only is not implemented by this worker", + )); + } + if let Some(strategy) = request.get("strategy") { + let strategy = strategy + .as_str() + .ok_or_else(|| Reject::bad_request("strategy must be a string"))?; + if !matches!(strategy, "auto" | "single") { + return Err(Reject::bad_request(format!( + "strategy {strategy:?} is not implemented by this worker" + ))); + } + } + let options: Vec = match kind { + Kind::Noul => vec!["false".into(), "true".into()], + Kind::Score => (0..6).map(|i| i.to_string()).collect(), + Kind::Choice => match request.get("options") { + Some(Value::Array(list)) => { + let mut opts = Vec::with_capacity(list.len()); + for (i, v) in list.iter().enumerate() { + match v.as_str() { + Some(o) => opts.push(o.to_owned()), + None => { + return Err(Reject::bad_request(format!( + "options[{i}] must be a string" + ))); + } + } + } + opts + } + Some(Value::Null) | None => Vec::new(), + Some(_) => { + return Err(Reject::bad_request("options must be a list of strings")); + } + }, + }; + if kind == Kind::Choice && !(2..=labels.len()).contains(&options.len()) { + return Err(Reject::bad_request(format!( + "choice needs 2-{} options, got {}", + labels.len(), + options.len() + ))); + } + let state = request + .get("state") + .cloned() + .unwrap_or(Value::String(String::new())); + let (parts, images) = parts(&state); + let lines: Vec = match kind { + Kind::Choice => options + .iter() + .enumerate() + .map(|(i, o)| format!("{}) {o}", labels[i])) + .collect(), + _ => options.clone(), + }; + let mut prompt = format!("[kind] {}\n[state] ", kind.as_str()); + for part in &parts { + match part { + Part::Text(t) => prompt.push_str(t), + Part::Image(_) => prompt.push_str(IMAGE_PLACEHOLDER), + } + } + prompt.push_str(&format!( + "\n[question] {question}\n[options]\n{}\n[decision]:", + lines.join("\n") + )); + Ok(Compiled { + kind, + parts, + question, + options, + prompt, + images, + }) +} + +/// The prompt tail after the last image placeholder, rebuilt exactly as +/// `compile` would have it. The L1 hit path tokenizes only this suffix; +/// `<|vision_end|>` is a special token (always atomic to BPE), so the split +/// concatenation equals a whole-prompt tokenization bit for bit. +/// None when the state carries no image (=> no prefix structure). +pub fn tail_after_last_image(compiled: &Compiled, labels: &[String]) -> Option { + let idx = compiled + .parts + .iter() + .rposition(|p| matches!(p, Part::Image(_)))?; + let mut tail = String::new(); + for part in &compiled.parts[idx + 1..] { + match part { + Part::Text(t) => tail.push_str(t), + Part::Image(_) => return None, + } + } + let lines: Vec = match compiled.kind { + Kind::Choice => compiled + .options + .iter() + .enumerate() + .map(|(i, o)| format!("{}) {o}", labels[i])) + .collect(), + _ => compiled.options.clone(), + }; + tail.push_str(&format!( + "\n[question] {}\n[options]\n{}\n[decision]:", + compiled.question, + lines.join("\n") + )); + Some(tail) +} + +/// The /v1/systemone answer, field-aligned with the official /v1/decide body. +pub fn answer( + kind: Kind, + options: &[String], + probabilities: &[f64], + input_tokens: usize, + elapsed_seconds: f64, +) -> Value { + let k = (1..probabilities.len()).fold(0, |m, i| { + if probabilities[i] > probabilities[m] { + i + } else { + m + } + }); + let adaptation = if kind != Kind::Choice || options.len() <= 16 { + "native" + } else { + "single" + }; + json!({ + "kind": kind.as_str(), + "effective_kind": kind.as_str(), + "options": options, + "probabilities": probabilities, + "choice_index": k, + "choice": options[k], + "adaptation": adaptation, + "protocol": PROTOCOL, + "model": MODEL_ID, + "usage": { + "prompt_tokens": input_tokens, + "completion_tokens": 1, + "total_tokens": input_tokens + 1, + }, + "elapsed_seconds": elapsed_seconds, + "num_model_requests": 1, + }) +} diff --git a/src/models/jev_vl/native/src/engine.rs b/src/models/jev_vl/native/src/engine.rs new file mode 100644 index 0000000..04d5da7 --- /dev/null +++ b/src/models/jev_vl/native/src/engine.rs @@ -0,0 +1,84 @@ +//! Assemble the processor and model executor from one pinned JEV-27B-VL export. + +use std::path::Path; +use std::sync::Arc; + +use anyhow::{Context, Result, ensure}; +use omni_runtime::SerialScheduler; +use serde_json::Value; + +use crate::caches::{CacheCfg, Caches}; +use crate::executor::{Executor, LabelHead}; +use crate::processing::Processor; + +pub struct Engine { + pub manifest: Value, + pub head: Arc, + pub processor: Processor, + pub scheduler: SerialScheduler, + pub executor: Executor, + /// Shared caches: processors/executors consult them; the HTTP layer + /// exposes its counters on /v1/cache/stats. + pub caches: Arc, +} + +impl Engine { + pub async fn load(dir: &Path, library: &Path) -> Result { + let manifest: Value = serde_json::from_slice( + &std::fs::read(dir.join("jev_vl_export.json")) + .context("export the merged checkpoint; see recipe/jev_vl/export_merged.py")?, + )?; + ensure!( + manifest["format"] == "jev-vl-text-merged/1" + && manifest["model_id"] == crate::contract::MODEL_ID + && manifest["protocol"] == crate::contract::PROTOCOL, + "expected a pinned JEV-27B-VL text export" + ); + let ranges: Vec = serde_json::from_value(manifest["slots"]["ranges"]["noul"].clone())?; + let score: Vec = serde_json::from_value(manifest["slots"]["ranges"]["score"].clone())?; + let choice: Vec = + serde_json::from_value(manifest["slots"]["ranges"]["choice"].clone())?; + ensure!( + ranges == [0, 2] && score == [2, 8] && choice == [8, 24], + "unexpected verbalizer slot ranges" + ); + ensure!( + manifest["verbalizer_ids"].as_array().map(Vec::len) == Some(24) + && manifest["verbalizer_bias"].as_array().map(Vec::len) == Some(24), + "24-slot verbalizer head" + ); + let labels: Vec = serde_json::from_value(manifest["labels"].clone())?; + let label_ids: Vec = serde_json::from_value(manifest["label_ids"].clone())?; + ensure!( + labels.len() == label_ids.len() && !labels.is_empty(), + "label table" + ); + for t in ["noul", "score", "choice"] { + let temp = manifest["temperatures"][t] + .as_f64() + .context("temperatures")?; + ensure!(temp.is_finite() && temp > 0.0, "invalid temperature"); + } + let max_length = manifest["max_length"].as_u64().context("max_length")? as usize; + ensure!( + (1..=32768).contains(&max_length), + "max_length must be within 1..=32768" + ); + let head = Arc::new(LabelHead::load(dir, &manifest)?); + let caches = Caches::new(CacheCfg::from_env()); + let model_index_hash = manifest["pins"]["model_index_sha256"] + .as_str() + .context("model_index_sha256")? + .to_owned(); + let processor = Processor::load(dir, labels, max_length, model_index_hash, caches.clone())?; + let executor = Executor::load(dir, library, head.clone(), caches.clone()).await?; + Ok(Self { + manifest, + head, + processor, + scheduler: SerialScheduler::default(), + executor, + caches, + }) + } +} diff --git a/src/models/jev_vl/native/src/executor.rs b/src/models/jev_vl/native/src/executor.rs new file mode 100644 index 0000000..461df3a --- /dev/null +++ b/src/models/jev_vl/native/src/executor.rs @@ -0,0 +1,425 @@ +//! Verbalizer readout: merged label-head rows on the CPU in f64 plus slot bias +//! and per-kind calibration temperature. +//! +//! The official readout uses full-vocabulary logprobs; the per-kind softmax only +//! sees `(logprob(t) + bias_slot) / T_kind`, where `logprob = logit - logZ` shares +//! one constant per request — so softmax((logit + bias)/T) is identical. The +//! export stores the merged lm_head rows for every exported label id, hence one +//! dot product per candidate token instead of a vocabulary GEMM. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use omni_qwen3_5_native::inputs::MultimodalInput; + +use anyhow::{Context, Result, ensure}; +use omni_qwen3_5_native::model::{Config, Model}; +use omni_runtime::SerialScheduler; +use serde_json::Value; + +use crate::images::ImageAsset; + +/// Per-request readout setup: which label tokens, with which bias, divided by +/// which temperature, matching decision_head.json's slots + calibration.json. +#[derive(Clone)] +pub struct Readout { + pub token_ids: Vec, + pub bias: Vec, + pub temperature: f64, +} + +/// The merged lm_head rows for the exported label ids, row-major float32. +pub struct LabelHead { + width: usize, + rows: Vec, + index: HashMap, +} + +impl LabelHead { + pub fn load(dir: &Path, manifest: &Value) -> Result { + let cfg = Config::load(dir)?; + ensure!( + ( + cfg.hidden, + cfg.intermediate, + cfg.full_attention.len(), + cfg.heads, + cfg.kv_heads, + cfg.lin_k_heads, + cfg.lin_v_heads + ) == (5120, 17408, 64, 24, 4, 16, 48), + "expected the Qwen3.8-27B backbone dimensions" + ); + let spec = manifest["label_head"].as_object().context("label_head")?["file"] + .as_str() + .context("label_head.file")?; + let data = std::fs::read(dir.join(spec)).context("read label_head safetensors")?; + let st = safetensors::SafeTensors::deserialize(&data)?; + let rows = st.tensor("rows")?; + ensure!( + rows.dtype() == safetensors::Dtype::F32 + && rows.shape().len() == 2 + && rows.shape()[1] == cfg.hidden, + "label_head rows must be float32 [n, {}]", + cfg.hidden + ); + let ids_t = st.tensor("ids")?; + ensure!( + ids_t.dtype() == safetensors::Dtype::I64 && ids_t.shape() == [rows.shape()[0]], + "label_head ids must be int64 [n]" + ); + let rows_f: Vec = rows + .data() + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) + .collect(); + let ids: Vec = ids_t + .data() + .as_chunks::<8>() + .0 + .iter() + .map(|b| i64::from_le_bytes(*b)) + .collect(); + let exported: Vec = serde_json::from_value(manifest["label_head"]["ids"].clone())?; + ensure!( + ids == exported, + "label_head ids do not match the export manifest" + ); + ensure!( + rows_f.iter().all(|v| v.is_finite()), + "non-finite label head" + ); + Ok(Self { + width: cfg.hidden, + rows: rows_f, + index: ids + .iter() + .enumerate() + .map(|(i, &t)| (t as u32, i)) + .collect(), + }) + } + + /// Slot bias from the trained 24-slot head: [slot_lo, slot_hi), then zeros. + pub fn readout(&self, manifest: &Value, kind_token: &str, count: usize) -> Result { + let temps = manifest["temperatures"] + .as_object() + .context("temperatures")?; + let temperature = temps[kind_token].as_f64().context("temperature")?; + ensure!( + temperature.is_finite() && temperature > 0.0, + "invalid temperature" + ); + let bias: Vec = manifest["verbalizer_bias"] + .as_array() + .context("verbalizer_bias")? + .iter() + .map(|v| v.as_f64().context("numeric bias")) + .collect::>()?; + let verbalizer: Vec = serde_json::from_value(manifest["verbalizer_ids"].clone())?; + ensure!(bias.len() == 24 && verbalizer.len() == 24, "24-slot head"); + let label_ids: Vec = serde_json::from_value(manifest["label_ids"].clone())?; + let (token_ids, slot_bias): (Vec, Vec) = match kind_token { + "noul" => (verbalizer[0..2].to_vec(), bias[0..2].to_vec()), + "score" => (verbalizer[2..8].to_vec(), bias[2..8].to_vec()), + "choice" => { + ensure!(count <= label_ids.len(), "choice options above label table"); + let b: Vec = (0..count) + .map(|i| if i < 16 { bias[8 + i] } else { 0.0 }) + .collect(); + (label_ids[..count].to_vec(), b) + } + other => anyhow::bail!("unknown kind {other}"), + }; + ensure!(token_ids.len() == slot_bias.len(), "readout shape"); + for &t in &token_ids { + ensure!( + self.index.contains_key(&(t as u32)), + "token {t} missing from the label head" + ); + } + Ok(Readout { + token_ids: token_ids.iter().map(|&t| t as u32).collect(), + bias: slot_bias, + temperature, + }) + } +} + +impl LabelHead { + /// One candidate distribution: `(logit + bias) / T` per slot then a stable + /// softmax. `last` is the final-norm hidden state as float32. + pub fn probabilities(&self, last: &[f32], readout: &Readout) -> Result> { + ensure!(last.len() == self.width, "hidden width"); + let w = self.width; + let mut z = Vec::with_capacity(readout.token_ids.len()); + for (t, b) in readout.token_ids.iter().zip(&readout.bias) { + let i = *self + .index + .get(t) + .with_context(|| format!("token {t} missing from the label head"))?; + let row = &self.rows[i * w..(i + 1) * w]; + z.push( + (row.iter() + .zip(last) + .map(|(&w, &h)| w as f64 * h as f64) + .sum::() + + b) + / readout.temperature, + ); + } + let base = z.iter().copied().fold(f64::NEG_INFINITY, f64::max); + let wgt: Vec = z.iter().map(|&v| (v - base).exp()).collect(); + let total: f64 = wgt.iter().sum(); + ensure!(total > 0.0 && total.is_finite(), "degenerate readout"); + Ok(wgt.iter().map(|v| v / total).collect()) + } +} + +/// One image block in the expanded request (adapted rows borrowed from L2). +pub struct ImgBlock { + /// Global expanded-id row of the first pad. + pub start: usize, + /// Global expanded-id row one past the last pad. + pub end: usize, + /// Parsed adapter output (grid + rows), shared with the L2 cache. + pub asset: Arc, +} + +/// Full expanded ids + absolute positions + image blocks (global coordinates). +pub struct PreparedMm { + pub ids: Vec, + pub positions: [Vec; 3], + pub blocks: Vec, +} + +/// The cached-prefix continuation slice of one request. +pub struct ContinueMm { + /// Rows [p, T): (pads_end - p) pads, <|vision_end|>, fresh tail tokens. + pub ids: Vec, + /// Absolute positions for those rows. + pub positions: [Vec; 3], + /// The image rows within the slice (the last block's tail), if any. + pub block: Option, + /// L3 device state of the prefix. + pub state: Arc, +} + +/// The image rows of one continuation slice: an offset range of one asset's rows. +pub struct SuffixBlock { + pub asset: Arc, + /// First adapted-row index feeding the slice (= p - pads_start). + pub row_offset: usize, + /// Number of rows = number of pads leading the slice. + pub rows: usize, +} + +/// The inputs and cached state needed to execute one request. +pub enum MmPlan { + /// Text-only prompt. + Text { ids: Vec }, + /// Full multimodal prefill. + Full(PreparedMm), + /// Structure recorded but its device state is not built yet: run the capture + /// phase over rows [0, p) followed by the continuation over [p, T), then + /// publish the state so later requests go straight to `Continue`. + Populate { + mm: PreparedMm, + record: Arc, + }, + /// Cached-prefix hit: only the continuation slice runs. + Continue(ContinueMm), +} + +pub struct Executor { + model: Arc>, + head: Arc, + caches: Arc, +} + +impl Executor { + pub(crate) async fn load( + dir: &Path, + library: &Path, + head: Arc, + caches: Arc, + ) -> Result { + let (d, lib) = (dir.to_path_buf(), library.to_path_buf()); + let model = tokio::task::spawn_blocking(move || Model::load(&d, &lib)).await??; + Ok(Self { + model: Arc::new(Mutex::new(model)), + head, + caches, + }) + } + + /// One forward, then `(logit + bias) / T` per slot and a stable softmax, + /// matching serve_decide.py's s1_pass ordering. + pub async fn execute( + &self, + scheduler: &SerialScheduler, + plan: MmPlan, + readout: Readout, + ) -> Result> { + let model = self.model.clone(); + let head = self.head.clone(); + let caches = self.caches.clone(); + scheduler + .run(move || { + let mut model = model + .lock() + .map_err(|_| anyhow::anyhow!("poisoned model"))?; + let result = (|| { + let last = match &plan { + MmPlan::Text { ids } => model.forward(ids)?, + MmPlan::Full(mm) => { + let one = slice_mm(mm, 0, mm.ids.len()); + let input = one.input(); + model.forward_multimodal(&input)? + } + MmPlan::Populate { mm, record } => { + let p = record.meta.p; + let (populated, state) = match model.alloc_prefix(p) { + Ok(mut state) => { + let result = (|| -> Result> { + let pf = slice_mm(mm, 0, p); + model + .forward_multimodal_capture(&pf.input(), &mut state)?; + let cf = slice_mm(mm, p, mm.ids.len()); + model.forward_multimodal_continue(&cf.input(), &state) + })(); + // The captured buffers must stay alive until this drain, + // including a failed capture or continuation. + model.synchronize()?; + (result, Some(state)) + } + Err(e) => (Err(e), None), + }; + match populated { + Ok(last) => { + caches.record_publish_state(record.key, state.unwrap()); + caches.l3_populate(); + last + } + Err(e) => { + // Device allocation or capture failed: serve the + // request on the proven one-shot path instead. + eprintln!("prefix populate failed; one-shot fallback: {e:#}"); + caches.l3_fallback_full(); + let one = slice_mm(mm, 0, mm.ids.len()); + let input = one.input(); + model.forward_multimodal(&input)? + } + } + } + MmPlan::Continue(c) => { + let indices: Vec = match &c.block { + Some(b) => (0..b.rows).collect(), + None => Vec::new(), + }; + let embeddings: &[half::bf16] = match &c.block { + Some(b) => { + &b.asset.embeddings + [b.row_offset * 5120..(b.row_offset + b.rows) * 5120] + } + None => &[], + }; + let input = MultimodalInput { + token_ids: &c.ids, + image_token_indices: &indices, + image_embeddings: embeddings, + position_ids: [ + c.positions[0].as_slice(), + c.positions[1].as_slice(), + c.positions[2].as_slice(), + ], + }; + model.forward_multimodal_continue(&input, &c.state)? + } + }; + head.probabilities(&last, &readout) + })(); + // Keep admission until all queued device work has drained, including errors. + model.synchronize()?; + result + }) + .await + } +} + +/// A slice of one request's expanded data for rows [from, to) in slice-local +/// coordinates: borrowed ids/positions, local pad indices, and the adapted rows +/// those pads need — borrowed straight out of the L2 asset when they form one +/// contiguous run (the usual single-image case), else concatenated once. +struct SlicedMm<'a> { + ids: &'a [u32], + positions: [&'a [i64]; 3], + indices: Vec, + borrowed: Option<&'a [half::bf16]>, + owned: Option>, +} + +impl<'a> SlicedMm<'a> { + fn input(&'a self) -> MultimodalInput<'a> { + MultimodalInput { + token_ids: self.ids, + image_token_indices: &self.indices, + image_embeddings: self.borrowed.or(self.owned.as_deref()).unwrap_or(&[]), + position_ids: self.positions, + } + } +} + +fn slice_mm<'a>(mm: &'a PreparedMm, from: usize, to: usize) -> SlicedMm<'a> { + let indices: Vec = mm + .blocks + .iter() + .flat_map(|b| (b.start.max(from)..b.end.min(to)).map(move |g| g - from)) + .collect(); + let overlapping: Vec<&ImgBlock> = mm + .blocks + .iter() + .filter(|b| b.end > from && b.start < to) + .collect(); + let (borrowed, owned) = match overlapping.as_slice() { + [] => (None, None), + [one] => { + let rows = (one.start.max(from) - one.start)..(one.end.min(to) - one.start); + ( + Some(&one.asset.embeddings[rows.start * 5120..rows.end * 5120]), + None, + ) + } + _ => ( + None, + Some( + overlapping + .iter() + .flat_map(|b| { + let rows = (b.start.max(from) - b.start)..(b.end.min(to) - b.start); + &b.asset.embeddings[rows.start * 5120..rows.end * 5120] + }) + .copied() + .collect::>(), + ), + ), + }; + SlicedMm { + ids: &mm.ids[from..to], + positions: [ + &mm.positions[0][from..to], + &mm.positions[1][from..to], + &mm.positions[2][from..to], + ], + indices, + borrowed, + owned, + } +} + +#[cfg(test)] +#[path = "../../../../../tests/jev_vl/readout.rs"] +mod readout_tests; diff --git a/src/models/jev_vl/native/src/images.rs b/src/models/jev_vl/native/src/images.rs new file mode 100644 index 0000000..d108954 --- /dev/null +++ b/src/models/jev_vl/native/src/images.rs @@ -0,0 +1,191 @@ +//! Adapted image assets (pre-encoded by recipe/jev_vl/preencode.py) and the +//! language-side position expansion for Qwen mrope, ported from HF transformers +//! modeling_qwen3_5 get_rope_index / get_vision_position_ids. + +use std::path::Path; +use std::sync::Arc; + +use anyhow::{Context, Result, ensure}; + +/// One adapted image: merger output rows in placeholder order + the patch grid. +pub struct ImageAsset { + pub grid_thw: [i64; 3], + pub embeddings: Vec, +} + +impl ImageAsset { + pub fn n_tokens(&self) -> usize { + let [t, h, w] = self.grid_thw; + (t * h * w / 4) as usize + } + + pub fn load(dir: &Path, url_hash: &str, model_index_hash: &str) -> Result> { + let grid: serde_json::Value = serde_json::from_slice( + &std::fs::read(dir.join("grid.json")).context("imgcache asset grid.json")?, + )?; + ensure!( + grid["url_sha256"].as_str() == Some(url_hash) + && grid["model_index_sha256"].as_str() == Some(model_index_hash), + "image asset source does not match the request and model index" + ); + let thw: [i64; 3] = serde_json::from_value(grid["grid_thw"].clone())?; + ensure!( + thw[0] == 1 && thw[1] > 0 && thw[2] > 0 && thw[1] % 2 == 0 && thw[2] % 2 == 0, + "expected one image with an even, positive spatial grid" + ); + let asset_n = thw[1] + .checked_mul(thw[2]) + .and_then(|v| usize::try_from(v / 4).ok()) + .context("image grid is too large")?; + ensure!(asset_n <= 32768, "image exceeds the supported token limit"); + ensure!( + grid["n_tokens"].as_u64() == Some(asset_n as u64), + "grid.json n_tokens mismatch" + ); + let data = std::fs::read(dir.join("emb.safetensors")).context("imgcache asset emb")?; + let st = safetensors::SafeTensors::deserialize(&data)?; + let rows = st.tensor("rows")?; + ensure!( + rows.dtype() == safetensors::Dtype::BF16 && rows.shape() == [asset_n, 5120], + "image rows vs grid mismatch" + ); + let embeddings: Vec = rows + .data() + .as_chunks::<2>() + .0 + .iter() + .map(|b| half::bf16::from_le_bytes(*b)) + .collect(); + ensure!( + embeddings.iter().all(|x| x.is_finite()), + "non-finite image embedding" + ); + Ok(Arc::new(Self { + grid_thw: [thw[0], thw[1], thw[2]], + embeddings, + })) + } +} + +/// The token-expanded multimodal prompt for the language-side boundary. +pub struct Expanded { + pub ids: Vec, + pub positions: [Vec; 3], + /// Each image's pad run in expanded-id coordinates plus its position base + /// and mrope advance, used to split a cached prefix from the question. + pub blocks: Vec, +} + +/// One image block's pad run in expanded coordinates. +#[derive(Clone, Debug)] +pub struct ImageBlock { + /// ids row of the block's first pad. + pub start: usize, + /// ids row one past the block's last pad. + pub end: usize, + /// The meshgrid position base of the block (equals `start` for the first + /// image; earlier images' mrope advances make it smaller for later ones). + pub base: i64, + /// mrope advance after the image: max(grid_h, grid_w) / 2. + pub advance: i64, + /// Index of the image in the request's image list. + pub asset: usize, +} + +/// mrope meshgrid triples for the image rows [from, from + rows) of a grid, +/// offset by `base` — a slice of what `expand` would build for the full block. +pub fn meshgrid_positions(grid: [i64; 3], base: i64, from: usize, rows: usize) -> [Vec; 3] { + let [gt, gh, gw] = grid; + let (lh, lw) = ((gh / 2) as usize, (gw / 2) as usize); + let _ = gt; + let mut out = [ + Vec::with_capacity(rows), + Vec::with_capacity(rows), + Vec::with_capacity(rows), + ]; + for l in from..from + rows { + let w = l % lw; + let h = (l / lw) % lh; + let t = l / (lh * lw); + out[0].push(base + t as i64); + out[1].push(base + h as i64); + out[2].push(base + w as i64); + } + out +} + +/// Tokens for one image block: the raw prompt text carries one literal +/// `<|vision_start|><|image_pad|><|vision_end|>` per image; the token sequence +/// repeats image_pad n = prod(grid)/merge² (merge=2) times. The three mrope +/// axes follow HF's meshgrid (t outer, h mid, w fastest), each offset by the +/// running start position; the next text token continues at +/// start + max(grid_h, grid_w) / merge. +pub fn expand(text_ids: &[u32], image_pad: u32, assets: &[Arc]) -> Result { + let marks: Vec = text_ids + .iter() + .enumerate() + .filter_map(|(i, &id)| (id == image_pad).then_some(i)) + .collect(); + ensure!( + marks.len() == assets.len(), + "prompt has {} image placeholders but the request carries {} images", + marks.len(), + assets.len() + ); + let mut ids: Vec = Vec::new(); + let mut positions: [Vec; 3] = [Vec::new(), Vec::new(), Vec::new()]; + let mut blocks: Vec = Vec::with_capacity(marks.len()); + let mut cursor = 0usize; + let mut current = 0i64; + for (k, &mark) in marks.iter().enumerate() { + for rel in 0..(mark - cursor) as i64 { + positions[0].push(current + rel); + positions[1].push(current + rel); + positions[2].push(current + rel); + } + ids.extend_from_slice(&text_ids[cursor..mark]); + current += (mark - cursor) as i64; + cursor = mark + 1; + let [gt, gh, gw] = assets[k].grid_thw; + ensure!( + gt >= 1 + && gh % 2 == 0 + && gw % 2 == 0 + && gt * gh * gw / 4 == assets[k].n_tokens() as i64, + "grid does not match the adapted rows" + ); + let (lg_t, lg_h, lg_w) = (gt, gh / 2, gw / 2); // temp_merge=1, spatial_merge=2 + let start = ids.len(); + for _ in 0..assets[k].n_tokens() { + ids.push(image_pad); + } + blocks.push(ImageBlock { + start, + end: ids.len(), + base: current, + advance: gh.max(gw) / 2, + asset: k, + }); + for ti in 0..lg_t { + for hi in 0..lg_h { + for wi in 0..lg_w { + positions[0].push(current + ti); + positions[1].push(current + hi); + positions[2].push(current + wi); + } + } + } + current += gh.max(gw) / 2; + } + for rel in 0..(text_ids.len() - cursor) as i64 { + positions[0].push(current + rel); + positions[1].push(current + rel); + positions[2].push(current + rel); + } + ids.extend_from_slice(&text_ids[cursor..]); + Ok(Expanded { + ids, + positions, + blocks, + }) +} diff --git a/src/models/jev_vl/native/src/lib.rs b/src/models/jev_vl/native/src/lib.rs new file mode 100644 index 0000000..f34f4f8 --- /dev/null +++ b/src/models/jev_vl/native/src/lib.rs @@ -0,0 +1,11 @@ +//! Native Rust/CUDA System-1 worker for autotrust/JEV-27B-VL: the merged text +//! backbone runs one prefill per request on the shared Qwen language kernels, and +//! the trained verbalizer readout (24-slot head + calibration) turns the last +//! hidden state into the official per-kind probability distribution. + +pub mod caches; +pub mod contract; +pub mod engine; +pub mod executor; +pub mod images; +pub mod processing; diff --git a/src/models/jev_vl/native/src/main.rs b/src/models/jev_vl/native/src/main.rs new file mode 100644 index 0000000..09117f3 --- /dev/null +++ b/src/models/jev_vl/native/src/main.rs @@ -0,0 +1,167 @@ +//! JEV_VL_MODEL= omni-jev-vl-native + +use std::sync::Arc; + +use anyhow::{Context, Result, ensure}; +use axum::{ + Json, Router, + body::Bytes, + extract::{DefaultBodyLimit, State}, + http::{HeaderMap, StatusCode, Uri}, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use omni_jev_vl_native::contract::{MODEL_ID, Reject}; +use omni_jev_vl_native::engine::Engine; +use omni_qwen3_5_native::cuda; +use serde_json::{Value, json}; + +const WARMUP: &[u8] = br#"{"kind":"choice","state":"Dialog: Update installed.","question":"Close it?","options":["OK","Wait"]}"#; + +fn reject(r: Reject) -> Response { + ( + StatusCode::from_u16(r.status).unwrap_or(StatusCode::BAD_REQUEST), + Json(r.body), + ) + .into_response() +} + +async fn decide(engine: &Engine, raw: &[u8]) -> Response { + let prepared = match engine + .processor + .prepare(&engine.head, &engine.manifest, raw) + { + Ok(prepared) => prepared, + Err(r) => return reject(r), + }; + let cache_note = prepared.cache_note.clone(); + let result = async { + let probs = engine + .executor + .execute(&engine.scheduler, prepared.plan, prepared.readout) + .await?; + anyhow::Ok(prepared.context.finish(probs)) + } + .await; + match result { + Ok(body) => { + let mut resp = Json(body).into_response(); + if !cache_note.is_empty() { + // Expose cache decisions without changing the response schema. + resp.headers_mut().insert( + "x-jev-cache", + axum::http::HeaderValue::from_str(&cache_note) + .unwrap_or(axum::http::HeaderValue::from_static("l1=?,l3=?,p=?")), + ); + } + resp + } + Err(e) => { + eprintln!("inference failed: {e:#}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"error": "model inference failed"})), + ) + .into_response() + } + } +} + +async fn systemone(State(engine): State>, headers: HeaderMap, body: Bytes) -> Response { + // Upstream pydantic validates the raw body even for a non-JSON content type; + // mirroring it keeps the wrong-content-type probe's 400 (not 415) semantics. + let ctype = headers + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .split(';') + .next() + .unwrap_or("") + .trim(); + if !ctype.eq_ignore_ascii_case("application/json") { + return reject(Reject::validation( + format!( + "1 validation error:\n {{'type': 'model_attributes_type', 'loc': 'body', 'msg': 'Input should be a valid dictionary or object to extract fields from', 'input': ''}}", + body.len() + ), + "body", + )); + } + decide(&engine, &body).await +} + +async fn fallback(uri: Uri, body: Bytes) -> Response { + let model = serde_json::from_slice::(&body) + .ok() + .and_then(|v| v.get("model").and_then(Value::as_str).map(str::to_owned)); + let _ = uri; + reject(Reject::unknown_model( + model.as_deref().unwrap_or("no-such-model"), + )) +} + +/// Cumulative cache counters and resident sizes, separate from decision responses. +async fn cache_stats(State(engine): State>) -> Response { + let s = engine.caches.snapshot(); + let cfg = &engine.caches.cfg; + Json(json!({ + "enabled": cfg.enabled, + "l1": cfg.l1, "l2": cfg.l2, "l3": cfg.l3, + "counters": { + "l1_hit": s.l1_hit, "l1_miss": s.l1_miss, + "l2_hit": s.l2_hit, "l2_miss": s.l2_miss, + "l3_hit": s.l3_hit, "l3_miss": s.l3_miss, + "l3_populate": s.l3_populate, "l3_fallback_full": s.l3_fallback_full, + }, + "records": { + "l1_records": s.l1_records, "l1_bytes": s.l1_bytes, + "l2_records": s.l2_records, "l2_bytes": s.l2_bytes, + "l3_records": s.l3_records, "l3_bytes": s.l3_bytes, + }, + "budgets": { + "l1_max": cfg.l1_max, "l2_bytes": cfg.l2_bytes, "l3_bytes": cfg.l3_bytes, + }, + })) + .into_response() +} + +async fn cache_reset(State(engine): State>) -> Response { + engine.caches.reset_stats(); + Json(json!({"reset": true})).into_response() +} + +#[tokio::main] +async fn main() -> Result<()> { + let model = std::env::var_os("JEV_VL_MODEL").context("set JEV_VL_MODEL")?; + let library = std::env::var_os("JEV_VL_CUDA_LIB") + .map(Into::into) + .map_or_else(cuda::default_library, Ok)?; + let engine = Arc::new(Engine::load(model.as_ref(), &library).await?); + ensure!( + decide(&engine, WARMUP).await.status() == StatusCode::OK, + "warmup failed" + ); + let host = std::env::var("JEV_VL_HOST").unwrap_or_else(|_| "127.0.0.1".into()); + let port: u16 = std::env::var("JEV_VL_PORT") + .map_or(Ok(8001), |v| v.parse()) + .context("JEV_VL_PORT")?; + let app = Router::new() + .route( + "/health", + get(|| async { Json(json!({"status": "ready", "model": MODEL_ID})) }), + ) + .route("/v1/systemone", post(systemone)) + .route("/v1/cache/stats", get(cache_stats)) + .route("/v1/cache/reset", post(cache_reset)) + .fallback(fallback) + .layer(DefaultBodyLimit::max(4 << 20)) + .with_state(engine); + let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; + println!("listening on {host}:{port}"); + axum::serve(listener, app).await?; + Ok(()) +} + +#[cfg(test)] +#[path = "../../../../../tests/jev_vl/http.rs"] +mod http_tests; diff --git a/src/models/jev_vl/native/src/processing.rs b/src/models/jev_vl/native/src/processing.rs new file mode 100644 index 0000000..7ac5b55 --- /dev/null +++ b/src/models/jev_vl/native/src/processing.rs @@ -0,0 +1,419 @@ +//! Request validation, tokenization, and response assembly. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Instant; + +use anyhow::{Context, Result}; +use serde_json::Value; +use tokenizers::Tokenizer; + +use crate::caches::{Caches, L1Meta, structure_key}; +use crate::contract::{self, Compiled, Kind, Reject}; +use crate::executor::{ContinueMm, ImgBlock, MmPlan, PreparedMm, Readout, SuffixBlock}; +use crate::images::{self, ImageAsset}; + +pub struct Processor { + tokenizer: Tokenizer, + labels: Vec, + max_length: usize, + imgcache: Option, + model_index_hash: String, + image_pad: u32, + vision_end: u32, + caches: Arc, +} + +pub struct PreparedRequest { + pub plan: MmPlan, + pub readout: Readout, + pub context: ResponseContext, + /// x-jev-cache response-header marker: l1/l2/l3 hit flags + prefix length. + pub cache_note: String, +} + +/// The original request mapping, usage, and timing. +pub struct ResponseContext { + kind: Kind, + options: Vec, + input_tokens: usize, + start: Instant, +} + +impl Processor { + pub(crate) fn load( + dir: &Path, + labels: Vec, + max_length: usize, + model_index_hash: String, + caches: Arc, + ) -> Result { + let tokenizer = + Tokenizer::from_file(dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + let image_pad = tokenizer + .token_to_id("<|image_pad|>") + .context("tokenizer has no <|image_pad|>")?; + let vision_end = tokenizer + .token_to_id("<|vision_end|>") + .context("tokenizer has no <|vision_end|>")?; + let hub = std::env::var("JEV_VL_IMGCACHE") + .ok() + .map(PathBuf::from) + .or_else(|| Some(dir.join("imgcache"))) + .filter(|p| p.is_dir()); + Ok(Self { + tokenizer, + labels, + max_length, + imgcache: hub, + model_index_hash, + image_pad, + vision_end, + caches, + }) + } + + fn url_key(url: &str) -> String { + use sha2::Digest; + sha2::Sha256::digest(url.as_bytes()) + .iter() + .map(|b| format!("{b:02x}")) + .collect() + } + + /// Load a prepared image asset, reusing the bounded L2 cache when enabled. + /// With caches disabled, each request reads and parses the asset again. + fn load_asset(&self, url: &str) -> Result, Reject> { + let key = Self::url_key(url); + if let Some(hit) = self.caches.l2_get(&key) { + return Ok(hit); + } + let dir = match &self.imgcache { + Some(d) => d.join(&key), + None => { + return Err(Reject::bad_request( + "image input: no imgcache is configured for this worker", + )); + } + }; + let asset = ImageAsset::load(&dir, &key, &self.model_index_hash) + .map_err(|e| Reject::bad_request(format!("image input is not preencoded: {e}")))?; + self.caches.l2_insert(key, asset.clone()); + Ok(asset) + } + + fn tokenize(&self, text: &str) -> Result, Reject> { + Ok(self + .tokenizer + .encode(text, false) + .map_err(|e| Reject::bad_request(format!("tokenization failed: {e}")))? + .get_ids() + .to_vec()) + } + + /// Validate, compile the raw prompt, and tokenize; add_special_tokens=false, + /// like the official raw System-1 readout. Never truncates: oversize prompts + /// get the upstream context-length rejection. + pub fn prepare( + &self, + head: &crate::executor::LabelHead, + manifest: &Value, + raw: &[u8], + ) -> Result { + let start = Instant::now(); + let compiled = contract::compile(raw, &self.labels)?; + let readout = head + .readout(manifest, compiled.kind.as_str(), compiled.options.len()) + .map_err(|e| Reject::bad_request(format!("invalid readout setup: {e:#}")))?; + let cfg = &self.caches.cfg; + if compiled.images.is_empty() || !cfg.enabled { + // Text-only requests and disabled caches use a full forward. + let text_ids = self.tokenize(&compiled.prompt)?; + let (plan, input_tokens, note) = if compiled.images.is_empty() { + ( + MmPlan::Text { ids: text_ids }, + 0, + "l1=-,l3=-,p=-".to_string(), + ) + } else { + let assets: Vec> = compiled + .images + .iter() + .map(|url| self.load_asset(url)) + .collect::, Reject>>()?; + let e = images::expand(&text_ids, self.image_pad, &assets) + .map_err(|e| Reject::bad_request(format!("invalid image prompt: {e:#}")))?; + let n = e.ids.len(); + ( + MmPlan::Full(expanded_mm(e, &assets)), + n, + "l1=off,l3=off,p=-".to_string(), + ) + }; + let input_tokens = match &plan { + MmPlan::Text { ids } => ids.len(), + _ => input_tokens, + }; + return self.finish_c(input_tokens, plan, readout, &compiled, start, note); + } + self.prepare_cached(&compiled, readout, start) + } + + /// L1 structure record + L3 continuation planning for one-image requests. + /// The plan decides whether the model sees the full prompt (miss/populate) or + /// only the continuation slice (record state present). + fn prepare_cached( + &self, + compiled: &Compiled, + readout: Readout, + start: Instant, + ) -> Result { + let cfg = &self.caches.cfg; + let key = structure_key(compiled.kind.as_str(), &compiled.parts); + let single_image = compiled.images.len() == 1; + let tail = contract::tail_after_last_image(compiled, &self.labels); + if let (Some(record), true) = (self.caches.record_get(key), single_image && tail.is_some()) + { + // Scalars out of the record up front: the borrow ends so the record + // can move into the plan arms below. + let meta = &record.meta; + let (p, pads_start, pads_end, base_pad, advance) = ( + meta.p, + meta.pads_start, + meta.pads_end, + meta.base_pad, + meta.advance, + ); + let asset = self.load_asset(&compiled.images[0])?; + // Suffix piece: pads(E-P) + <|vision_end|> + fresh tail tokenization. + let tail = tail.unwrap(); + let tail_ids = self.tokenize(&tail)?; + if tail_ids.contains(&self.image_pad) { + return Err(Reject::bad_request( + "invalid image prompt: unexpected image placeholder in suffix", + )); + } + let suffix_pads = pads_end - p; + let mut ids = Vec::with_capacity(suffix_pads + 1 + tail_ids.len()); + ids.extend(std::iter::repeat_n(self.image_pad, suffix_pads)); + ids.push(self.vision_end); + ids.extend_from_slice(&tail_ids); + let mut pos = + images::meshgrid_positions(asset.grid_thw, base_pad, p - pads_start, suffix_pads); + let after = base_pad + advance; + for k in 0..(ids.len() - suffix_pads) as i64 { + pos[0].push(after + k); + pos[1].push(after + k); + pos[2].push(after + k); + } + let input_tokens = p + ids.len(); + match (cfg.l3, record.state()) { + (true, Some(state)) => { + self.caches.l3_hit(); + let block = Some(SuffixBlock { + asset, + row_offset: p - pads_start, + rows: suffix_pads, + }); + return self.finish_c( + input_tokens, + MmPlan::Continue(ContinueMm { + ids, + positions: pos, + block, + state, + }), + readout, + compiled, + start, + format!("l1=hit,l3=hit,p={p}"), + ); + } + (true, None) => { + self.caches.l3_miss(); + let mm = rebuild_full(&ids, &pos, &record.meta, &asset); + return self.finish_c( + mm.ids.len(), + MmPlan::Populate { mm, record }, + readout, + compiled, + start, + format!("l1=hit,l3=miss,p={p}"), + ); + } + (false, _) => { + let mm = rebuild_full(&ids, &pos, &record.meta, &asset); + return self.finish_c( + mm.ids.len(), + MmPlan::Full(mm), + readout, + compiled, + start, + "l1=hit,l3=off,p=-".to_string(), + ); + } + } + } + // L1 miss: full expansion (and the cache record when the shape qualifies). + let text_ids = self.tokenize(&compiled.prompt)?; + let assets: Vec> = compiled + .images + .iter() + .map(|url| self.load_asset(url)) + .collect::, Reject>>()?; + let e = images::expand(&text_ids, self.image_pad, &assets) + .map_err(|e| Reject::bad_request(format!("invalid image prompt: {e:#}")))?; + let input_tokens = e.ids.len(); + let meta = (single_image && tail.is_some() && cfg.l1) + .then(|| prefix_meta(&e, input_tokens)) + .flatten(); + if let Some(meta) = meta { + let p = meta.p; + let record = self.caches.record_insert(key, meta); + if cfg.l3 { + self.caches.l3_miss(); + return self.finish_c( + input_tokens, + MmPlan::Populate { + mm: expanded_mm(e, &assets), + record, + }, + readout, + compiled, + start, + format!("l1=miss,l3=miss,p={p}"), + ); + } + return self.finish_c( + input_tokens, + MmPlan::Full(expanded_mm(e, &assets)), + readout, + compiled, + start, + "l1=miss,l3=off,p=-".to_string(), + ); + } + self.finish_c( + input_tokens, + MmPlan::Full(expanded_mm(e, &assets)), + readout, + compiled, + start, + "l1=miss,l3=off,p=-".to_string(), + ) + } + + fn finish_c( + &self, + input_tokens: usize, + plan: MmPlan, + readout: Readout, + compiled: &Compiled, + start: Instant, + cache_note: String, + ) -> Result { + let t = input_tokens; + if t > self.max_length { + return Err(Reject::bad_request(format!( + "This model's maximum context length is {} tokens. However, you requested 1 output tokens and your prompt contains at least {t} input tokens, for a total of at least {} tokens. Please reduce the length of the input prompt or the number of requested output tokens. (parameter=input_tokens, value={t})", + self.max_length, + t + 1 + ))); + } + Ok(PreparedRequest { + plan, + readout, + cache_note, + context: ResponseContext { + kind: compiled.kind, + options: compiled.options.clone(), + input_tokens: t, + start, + }, + }) + } +} + +/// Cache-anchor geometry for one expanded single-image prompt: prefix = rows +/// [0, floor(pads_end/64)*64) — a multiple of the 64-token GDN chunk, containing +/// the whole text head and all but a tail slice of the image block. +fn prefix_meta(e: &images::Expanded, input_tokens: usize) -> Option { + let block = e.blocks.last()?; + let p = block.end / 64 * 64; + if p < 64 || p < block.start || p >= input_tokens { + return None; + } + Some(L1Meta { + pads_start: block.start, + pads_end: block.end, + p, + base_pad: block.base, + advance: block.advance, + ids_prefix: e.ids[..p].to_vec(), + positions_prefix: [ + e.positions[0][..p].to_vec(), + e.positions[1][..p].to_vec(), + e.positions[2][..p].to_vec(), + ], + }) +} + +/// Full expanded ids/positions with image blocks borrowed from their assets. +fn expanded_mm(e: images::Expanded, assets: &[Arc]) -> PreparedMm { + let blocks = e + .blocks + .iter() + .map(|b| ImgBlock { + start: b.start, + end: b.end, + asset: assets[b.asset].clone(), + }) + .collect(); + PreparedMm { + ids: e.ids, + positions: e.positions, + blocks, + } +} + +/// Rebuild the full expanded ids/positions from the cached prefix plus the +/// freshly built suffix — exactly the bytes the miss path would have produced. +fn rebuild_full( + ids_suffix: &[u32], + pos_suffix: &[Vec; 3], + meta: &L1Meta, + asset: &Arc, +) -> PreparedMm { + let mut ids = meta.ids_prefix.clone(); + ids.extend_from_slice(ids_suffix); + let mut positions = meta.positions_prefix.clone(); + for a in 0..3 { + positions[a].extend_from_slice(&pos_suffix[a]); + } + let blocks = vec![ImgBlock { + start: meta.pads_start, + end: meta.pads_end, + asset: asset.clone(), + }]; + PreparedMm { + ids, + positions, + blocks, + } +} + +impl ResponseContext { + pub fn finish(self, probabilities: Vec) -> Value { + contract::answer( + self.kind, + &self.options, + &probabilities, + self.input_tokens, + self.start.elapsed().as_secs_f64(), + ) + } +} + +#[cfg(test)] +#[path = "../../../../../tests/jev_vl/processing.rs"] +mod tests; diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 3de0568..5e8141d 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -9,7 +9,7 @@ use std::sync::OnceLock; use anyhow::{Context, Result, bail, ensure}; /// `CS1_ABI_VERSION` in ops.h. -const ABI_VERSION: u32 = 5; +const ABI_VERSION: u32 = 6; pub const LIBRARY: &str = "libqwen3_5_cuda.so"; /// A `cudaStream_t`. @@ -73,6 +73,8 @@ api! { cs1_graph_destroy(exec: *mut c_void) -> c_int; cs1_upload(dst: *mut c_void, src: *const c_void, bytes: usize, stream: Stream) -> c_int; cs1_download(dst: *mut c_void, src: *const c_void, bytes: usize, stream: Stream) -> c_int; + cs1_copy_dd(dst: *mut c_void, src: *const c_void, bytes: usize, stream: Stream) -> c_int; + cs1_copy2d(dst: *mut c_void, dpitch: usize, src: *const c_void, spitch: usize, width: usize, height: usize, stream: Stream) -> c_int; cs1_embed(ids: *const i32, table: *const c_void, out: *mut c_void, t: c_int, d: c_int, stream: Stream) -> c_int; cs1_rms_norm( x: *const c_void, w: *const c_void, out: *mut c_void, rows: c_int, d: c_int, eps: f32, stream: Stream, @@ -98,6 +100,11 @@ api! { q: *const c_void, k: *const c_void, v: *const c_void, g: *const f32, beta: *const c_void, o: *mut c_void, workspace: *mut f32, t: c_int, h: c_int, hk: c_int, scale: f32, stream: Stream, ) -> c_int; + cs1_gdn_prefill_x( + q: *const c_void, k: *const c_void, v: *const c_void, g: *const f32, beta: *const c_void, o: *mut c_void, + workspace: *mut f32, t: c_int, h: c_int, hk: c_int, scale: f32, + s_in: *const c_void, s_out: *mut c_void, stream: Stream, + ) -> c_int; cs1_attn_prep( qg: *const c_void, kr: *const c_void, ld: c_int, qw: *const c_void, kw: *const c_void, cos: *const c_void, sin: *const c_void, q: *mut c_void, gate: *mut c_void, k: *mut c_void, t: c_int, hq: c_int, hk: c_int, @@ -111,6 +118,10 @@ api! { q: *const c_void, k: *const c_void, v: *const c_void, ldv: c_int, gate: *const c_void, out: *mut c_void, t: c_int, hq: c_int, hk: c_int, dh: c_int, scale: f32, stream: Stream, ) -> c_int; + cs1_attention_gated_prefix( + q: *const c_void, k: *const c_void, v: *const c_void, ldv: c_int, gate: *const c_void, + out: *mut c_void, t: c_int, hq: c_int, hk: c_int, dh: c_int, scale: f32, q_base: c_int, stream: Stream, + ) -> c_int; cs1_sigmoid_gate(x: *mut c_void, gate: *const c_void, n: usize, stream: Stream) -> c_int; cs1_silu_mul(gate_up: *const c_void, ld: c_int, out: *mut c_void, t: c_int, i: c_int, stream: Stream) -> c_int; cs1_gemm_create(workspace_bytes: usize) -> *mut c_void; @@ -258,6 +269,41 @@ pub unsafe fn download(dst: &mut [u8], src: *const c_void, stream: Stream) -> Re ) } +/// Queue a device-to-device copy of `bytes`; completion is only stream-ordered. +/// +/// # Safety +/// Both ranges of `bytes` must be valid device allocations, non-overlapping. +pub unsafe fn copy_dd( + dst: *mut c_void, + src: *const c_void, + bytes: usize, + stream: Stream, +) -> Result<()> { + check( + unsafe { (api().cs1_copy_dd)(dst, src, bytes, stream) }, + "device copy", + ) +} + +/// Queue a pitched device-to-device copy: `height` rows of `width` bytes. +/// +/// # Safety +/// `src`/`dst` must be device allocations with the given pitches and heights. +pub unsafe fn copy2d( + dst: *mut c_void, + dpitch: usize, + src: *const c_void, + spitch: usize, + width: usize, + height: usize, + stream: Stream, +) -> Result<()> { + check( + unsafe { (api().cs1_copy2d)(dst, dpitch, src, spitch, width, height, stream) }, + "pitched device copy", + ) +} + pub struct Graph { exec: *mut c_void, } diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 6d37cd2..304c308 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -12,6 +12,7 @@ use std::collections::{HashMap, VecDeque}; use std::ffi::c_void; use std::path::Path; +use std::sync::Arc; use anyhow::{Context, Result, bail, ensure}; use serde_json::Value as Json; @@ -553,8 +554,42 @@ impl Scratch { } } +/// The cached device-side state of one token prefix (rows `[0, len)`) for the R2d +/// hybrid path: per full-attention layer the post-prep K rows plus the pre-GEMM V +/// rows, per Gated DeltaNet layer the float32 recurrent state and the three +/// pre-conv projection columns feeding the conv window. States belong to the +/// prefix itself; a continuation seeds its own scratch from them, read-only. +pub struct PrefixState { + owner: Arc<()>, + initialized: bool, + /// Tokens this prefix covers; always a multiple of 64 (the GDN chunk length), + /// which also aligns the flash-attention key tiles of the one-shot pass. + len: usize, + /// Post-prep K [len, Hk*Dh] and raw V [len, Hk*Dh] rows per full-attention layer. + attn_kv: Vec<(DeviceBuffer, DeviceBuffer)>, + /// Float32 [H, K, V] recurrent states, one per Gated DeltaNet layer. + gdn_state: Vec, + /// Three pre-conv projection columns [3, gdn_in width], one per Gated DeltaNet layer. + conv_tail: Vec, + /// Total device bytes, for cache-budget accounting by the caller. + bytes: usize, +} + +impl PrefixState { + /// Number of tokens covered by this nonempty prefix. + pub fn token_count(&self) -> usize { + self.len + } + + /// Device allocation size used by cache-budget accounting. + pub fn bytes(&self) -> usize { + self.bytes + } +} + pub struct Model { pub cfg: Config, + prefix_owner: Arc<()>, _weights: Weights, embed: Tensor, final_norm: Tensor, @@ -647,6 +682,7 @@ impl Model { ensure!(!gemm.is_null(), "cuBLASLt setup failed"); let model = Self { cfg, + prefix_owner: Arc::new(()), _weights: weights, embed, final_norm, @@ -808,6 +844,36 @@ impl Model { let s = self.scratch.as_ref().unwrap(); self.upload_positions(s, input.position_ids)?; self.embed_tokens(s, input.token_ids)?; + self.overwrite_image_rows(s, input, 0)?; + self.run(s, &[t], true)?; + self.last_hidden(s, t) + } + + fn upload_positions(&self, s: &Scratch, positions: [&[i64]; 3]) -> Result<()> { + let (cos, sin) = rotary_tables( + positions, + self.cfg.rotary_half, + self.cfg.rope_theta, + self.cfg.mrope_section, + ); + // SAFETY: tables contain at most s.cap rows of rotary_half BF16 values. + unsafe { + cuda::upload(s.at(s.custom_cos), &cos, self.stream)?; + cuda::upload(s.at(s.custom_sin), &sin, self.stream)?; + } + Ok(()) + } + + /// Overwrite the residual rows of validated image placeholders with the adapted + /// image embeddings. `qb` is the absolute row base of the passed slice inside + /// the residual buffer (0 for a full prompt, the cached prefix length for a + /// continuation — indices are local to the slice in both cases). + fn overwrite_image_rows( + &self, + s: &Scratch, + input: &MultimodalInput<'_>, + qb: usize, + ) -> Result<()> { let bytes: Vec = input .image_embeddings .iter() @@ -822,46 +888,158 @@ impl Model { end += 1; } let row_bytes = self.cfg.hidden * BF16; - // SAFETY: validated indices lie in the t-row residual buffer, and - // features contain exactly one hidden-size BF16 row per placeholder. + // SAFETY: validated indices lie in the residual buffer, and features + // contain exactly one hidden-size BF16 row per placeholder. unsafe { cuda::upload( - s.at(s.res + indices[begin] * row_bytes), + s.at(s.res + (qb + indices[begin]) * row_bytes), &bytes[begin * row_bytes..end * row_bytes], self.stream, )?; } begin = end; } - self.run(s, &[t], true)?; - self.last_hidden(s, t) + Ok(()) } - fn upload_positions(&self, s: &Scratch, positions: [&[i64]; 3]) -> Result<()> { - let (cos, sin) = rotary_tables( - positions, - self.cfg.rotary_half, - self.cfg.rope_theta, - self.cfg.mrope_section, + /// Allocate an uninitialized prefix of `len` tokens (a multiple of 64). + /// Capture must finish successfully on this model before continuation. + pub fn alloc_prefix(&self, len: usize) -> Result { + ensure!( + len >= 64 && len.is_multiple_of(64), + "cached prefix length must be a positive multiple of 64" ); - // SAFETY: tables contain at most s.cap rows of rotary_half BF16 values. - unsafe { - cuda::upload(s.at(s.custom_cos), &cos, self.stream)?; - cuda::upload(s.at(s.custom_sin), &sin, self.stream)?; + ensure!( + len <= self.cfg.max_positions, + "cached prefix exceeds the configured maximum length" + ); + let kvrow = self.cfg.kv_heads * self.cfg.head_dim * BF16; + let w = Widths::of(&self.cfg); + let nfull = self.cfg.full_attention.iter().filter(|&&f| f).count(); + let ngdn = self.cfg.full_attention.len() - nfull; + let mut bytes = 0usize; + let mut attn_kv = Vec::with_capacity(nfull); + for _ in 0..nfull { + let k = DeviceBuffer::new(len * kvrow)?; + let v = DeviceBuffer::new(len * kvrow)?; + bytes += 2 * len * kvrow; + attn_kv.push((k, v)); + } + let state_floats = self.cfg.lin_v_heads * self.cfg.lin_k_dim * self.cfg.lin_v_dim; + let tail_bytes = 3 * w.gdn_in * BF16; + let mut gdn_state = Vec::with_capacity(ngdn); + let mut conv_tail = Vec::with_capacity(ngdn); + for _ in 0..ngdn { + gdn_state.push(DeviceBuffer::new(state_floats * F32)?); + conv_tail.push(DeviceBuffer::new(tail_bytes)?); + bytes += state_floats * F32 + tail_bytes; } + Ok(PrefixState { + owner: Arc::clone(&self.prefix_owner), + initialized: false, + len, + attn_kv, + gdn_state, + conv_tail, + bytes, + }) + } + + /// Run the rows `[0, state.len)` of one prompt, collecting the per-layer + /// prefix state (post-prep K/V columns, float32 GDN states, conv tails) into + /// `state`. Synchronize before marking the state ready for continuation. + /// The hidden states of these rows are computed but not returned. + pub fn forward_multimodal_capture( + &mut self, + input: &MultimodalInput<'_>, + state: &mut PrefixState, + ) -> Result<()> { + ensure!( + Arc::ptr_eq(&self.prefix_owner, &state.owner), + "prefix state belongs to a different model instance" + ); + let t = input.token_ids.len(); + ensure!( + t == state.len, + "capture input must cover the cached prefix exactly" + ); + state.initialized = false; + let image_token = self + .cfg + .image_token_id + .context("checkpoint has no image_token_id")?; + input.validate( + self.cfg.hidden, + self.embed.shape[0], + image_token, + self.cfg.max_positions, + )?; + self.prepare_scratch(t)?; + let s = self.scratch.as_ref().unwrap(); + self.upload_positions(s, input.position_ids)?; + self.embed_tokens(s, input.token_ids)?; + self.overwrite_image_rows(s, input, 0)?; + self.run_window(s, 0, t, true, Some(state), None)?; + self.synchronize()?; + state.initialized = true; Ok(()) } + /// Run rows `[prefix.len, prefix.len + input.len)` of one prompt seeded from a + /// captured prefix; returns the final-norm hidden state at the last position. + pub fn forward_multimodal_continue( + &mut self, + input: &MultimodalInput<'_>, + prefix: &PrefixState, + ) -> Result> { + ensure!( + Arc::ptr_eq(&self.prefix_owner, &prefix.owner), + "prefix state belongs to a different model instance" + ); + ensure!(prefix.initialized, "prefix state has not completed capture"); + let image_token = self + .cfg + .image_token_id + .context("checkpoint has no image_token_id")?; + input.validate( + self.cfg.hidden, + self.embed.shape[0], + image_token, + self.cfg.max_positions, + )?; + let rows = input.token_ids.len(); + let tend = prefix + .len + .checked_add(rows) + .context("cached prompt length overflow")?; + ensure!( + tend <= self.cfg.max_positions, + "cached prompt exceeds the configured maximum length" + ); + self.prepare_scratch(tend)?; + let s = self.scratch.as_ref().unwrap(); + self.upload_positions(s, input.position_ids)?; + self.embed_tokens_at(s, input.token_ids, prefix.len)?; + self.overwrite_image_rows(s, input, prefix.len)?; + self.run_window(s, prefix.len, tend, true, None, Some(prefix))?; + self.last_hidden(s, tend) + } + fn embed_tokens(&self, s: &Scratch, ids: &[u32]) -> Result<()> { + self.embed_tokens_at(s, ids, 0) + } + + /// Embed `ids` into the residual buffer rows starting at absolute row `qb`. + fn embed_tokens_at(&self, s: &Scratch, ids: &[u32], qb: usize) -> Result<()> { let ids32: Vec = ids.iter().flat_map(|&i| (i as i32).to_le_bytes()).collect(); - // SAFETY: IDs were checked against the vocabulary; scratch holds t rows. + // SAFETY: IDs were checked against the vocabulary; scratch holds qb + t rows. unsafe { cuda::upload(s.at(s.ids), &ids32, self.stream)?; check( (cuda::api().cs1_embed)( s.at(s.ids).cast(), self.embed.ptr, - s.at(s.res), + s.at(s.res + qb * self.cfg.hidden * BF16), ids.len() as i32, self.cfg.hidden as i32, self.stream, @@ -1111,4 +1289,304 @@ impl Model { } Ok(()) } + + /// The R2d prefix-cache pass: run rows `[qb, tend)` of one prompt with + /// explicit positions, touching per-layer buffers only for those rows. + /// + /// `capture` (qb must be 0): collect post-prep K and pre-GEMM V columns of each + /// full-attention layer plus the float32 GDN state and the three pre-conv + /// projection columns of each Gated DeltaNet layer. `prefix` (qb = prefix.len): + /// seed the window from a captured entry — the conv tail columns are copied + /// back into place before the convolution, and the prefix K/V columns before + /// the attention. Everything else is `run()` with shifted row bases, so a + /// captured+continued pass repeats the one-shot arithmetic token for token + /// (qb is a multiple of the 64-row GDN chunk; the only remaining divergence + /// is cuBLASLt's M-shaped plan choice, an f32-accumulation-order ulp). + fn run_window( + &self, + s: &Scratch, + qb: usize, + tend: usize, + custom_positions: bool, + capture: Option<&mut PrefixState>, + prefix: Option<&PrefixState>, + ) -> Result<()> { + ensure!( + qb <= tend && (qb == 0) == prefix.is_none() && (qb == 0 || capture.is_none()), + "prefix window shape" + ); + let cfg = &self.cfg; + let st = self.stream; + let rows = tend - qb; + let (ti, hi, eps) = (rows as i32, cfg.hidden as i32, cfg.eps); + let (kd, vd, hv) = (cfg.key_dim(), cfg.value_dim(), cfg.lin_v_heads); + let (hq, hk, hd) = (cfg.heads as i32, cfg.kv_heads as i32, cfg.head_dim as i32); + let w = Widths::of(cfg); + let p = |off: usize| s.at(off); + let (cos, sin) = if custom_positions { + (s.custom_cos, s.custom_sin) + } else { + (s.cos, s.sin) + }; + let hb = cfg.hidden * BF16; + let ob = cfg.heads * cfg.head_dim * BF16; // post-prep q/gate and ao row bytes + let kvb = cfg.kv_heads * cfg.head_dim * BF16; // k and v row bytes + let ldb = w.gdn_in; // gdn_in row, elements + let ldbb = ldb * BF16; + let ab = w.attn_in * BF16; // attn_in row bytes + let kb = kd * BF16; // linear q/k row bytes + let vb = vd * BF16; // linear v row bytes + let scale = (cfg.head_dim as f32).powf(-0.5); + let mut fa_i = 0usize; + let mut la_i = 0usize; + // SAFETY (every kernel call below): pointers are weights in the arena or + // scratch buffers laid out for tend tokens at the widths used here; prefix + // buffers hold exactly the shapes allocated by alloc_prefix(qb). + unsafe { + check( + (cuda::api().cs1_rms_norm)( + p(s.res + qb * hb), + self.layers[0].input_norm.ptr, + p(s.x + qb * hb), + ti, + hi, + eps, + st, + ), + "input norm", + )?; + } + for (i, layer) in self.layers.iter().enumerate() { + match &layer.mixer { + Mixer::Linear(la) => { + self.gemm(s, s.x + qb * hb, &la.in_proj, s.gdn_in + qb * ldbb, rows)?; + let ld = ldb as i32; + let z = s.gdn_in + w.conv * BF16; + let b = z + vd * BF16; + let a = b + hv * BF16; + let (s_in, s_out): (*const c_void, *mut c_void) = + match (&prefix, capture.as_deref()) { + (Some(pre), _) => { + (pre.gdn_state[la_i].at(0).cast_const(), std::ptr::null_mut()) + } + (None, Some(cap)) => (std::ptr::null(), cap.gdn_state[la_i].at(0)), + (None, None) => (std::ptr::null(), std::ptr::null_mut()), + }; + unsafe { + if let Some(cap) = capture.as_deref() { + // conv window tail: the last three pre-conv columns. + cuda::copy_dd( + cap.conv_tail[la_i].at(0), + p(s.gdn_in + (rows - 3) * ldbb), + 3 * ldbb, + st, + )?; + } + if qb > 0 { + cuda::copy_dd( + p(s.gdn_in + (qb - 3) * ldbb), + prefix.unwrap().conv_tail[la_i].at(0), + 3 * ldbb, + st, + )?; + } + // The conv kernel zero-pads its local first three rows. + // Include the restored history in its input and discard + // those three outputs so continuation rows see all taps. + let conv_start = qb.saturating_sub(3); + check( + (cuda::api().cs1_gdn_conv)( + p(s.gdn_in + conv_start * ldbb), + ld, + la.conv.ptr, + p(s.lq + conv_start * kb), + p(s.lk + conv_start * kb), + p(s.lv + conv_start * vb), + (tend - conv_start) as i32, + kd as i32, + vd as i32, + st, + ), + "gdn conv", + )?; + check( + (cuda::api().cs1_gdn_gates)( + p(b + qb * ldbb), + p(a + qb * ldbb), + ld, + la.a_log.ptr, + la.dt_bias.ptr, + p(s.beta + qb * hv * BF16), + p(s.g + qb * hv * F32).cast(), + ti, + hv as i32, + st, + ), + "gdn gates", + )?; + check( + (cuda::api().cs1_gdn_prefill_x)( + p(s.lq + qb * kb), + p(s.lk + qb * kb), + p(s.lv + qb * vb), + p(s.g + qb * hv * F32).cast(), + p(s.beta + qb * hv * BF16), + p(s.lo + qb * vb), + p(s.workspace).cast(), + ti, + hv as i32, + cfg.lin_k_heads as i32, + (cfg.lin_k_dim as f32).powf(-0.5), + s_in, + s_out, + st, + ), + "gdn prefill", + )?; + check( + (cuda::api().cs1_gated_rms_norm)( + p(s.lo + qb * vb), + p(z + qb * ldbb), + ld, + la.norm.ptr, + p(s.ln + qb * vb), + ti, + hv as i32, + cfg.lin_v_dim as i32, + eps, + st, + ), + "gated norm", + )?; + } + self.gemm(s, s.ln + qb * vb, &la.out, s.delta + qb * hb, rows)?; + la_i += 1; + } + Mixer::Full(fa) => { + self.gemm(s, s.x + qb * hb, &fa.qkv, s.attn_in + qb * ab, rows)?; + let k = s.attn_in + w.attn_q * BF16; + let v = k + kvb; + let (q_base, t_flash) = (qb as i32, tend as i32); + unsafe { + if qb > 0 { + let pre = prefix.unwrap(); + // Restore the prefix V columns and post-prep K rows in place; + // the windowed flash then sees the identical k/v layout. + cuda::copy2d(p(v), ab, pre.attn_kv[fa_i].1.at(0), kvb, kvb, qb, st)?; + cuda::copy_dd(p(s.ak), pre.attn_kv[fa_i].0.at(0), qb * kvb, st)?; + } + check( + (cuda::api().cs1_attn_prep)( + p(s.attn_in + qb * ab), + p(k + qb * ab), + w.attn_in as i32, + fa.q_norm.ptr, + fa.k_norm.ptr, + p(cos), + p(sin), + p(s.aq + qb * ob), + p(s.agate + qb * ob), + p(s.ak + qb * kvb), + ti, + hq, + hk, + hd, + cfg.rotary_half as i32, + eps, + st, + ), + "attention prep", + )?; + if let Some(cap) = capture.as_deref() { + // K must be post-prep (normed + rotary); V is the raw column. + cuda::copy_dd(cap.attn_kv[fa_i].0.at(0), p(s.ak), tend * kvb, st)?; + cuda::copy2d(cap.attn_kv[fa_i].1.at(0), kvb, p(v), ab, kvb, tend, st)?; + } + check( + (cuda::api().cs1_attention_gated_prefix)( + p(s.aq), + p(s.ak), + p(v), + w.attn_in as i32, + p(s.agate), + p(s.ao), + t_flash, + hq, + hk, + hd, + scale, + q_base, + st, + ), + "gated attention", + )?; + } + self.gemm(s, s.ao + qb * ob, &fa.o, s.delta + qb * hb, rows)?; + fa_i += 1; + } + } + unsafe { + check( + (cuda::api().cs1_add_rms_norm)( + p(s.res + qb * hb), + p(s.delta + qb * hb), + layer.post_norm.ptr, + p(s.x + qb * hb), + ti, + hi, + eps, + st, + ), + "post-attention norm", + )?; + } + self.gemm( + s, + s.x + qb * hb, + &layer.gate_up, + s.gate_up + qb * 2 * cfg.intermediate * BF16, + rows, + )?; + unsafe { + check( + (cuda::api().cs1_silu_mul)( + p(s.gate_up + qb * 2 * cfg.intermediate * BF16), + (2 * cfg.intermediate) as i32, + p(s.act + qb * cfg.intermediate * BF16), + ti, + cfg.intermediate as i32, + st, + ), + "silu mul", + )?; + } + self.gemm( + s, + s.act + qb * cfg.intermediate * BF16, + &layer.down, + s.delta + qb * hb, + rows, + )?; + let next = self + .layers + .get(i + 1) + .map_or(&self.final_norm, |l| &l.input_norm); + unsafe { + check( + (cuda::api().cs1_add_rms_norm)( + p(s.res + qb * hb), + p(s.delta + qb * hb), + next.ptr, + p(s.x + qb * hb), + ti, + hi, + eps, + st, + ), + "input norm", + )?; + } + } + Ok(()) + } } diff --git a/tests/benchmarks/test_jev_vl_preencode.py b/tests/benchmarks/test_jev_vl_preencode.py new file mode 100644 index 0000000..b80c010 --- /dev/null +++ b/tests/benchmarks/test_jev_vl_preencode.py @@ -0,0 +1,186 @@ +"""Check image selection and recoverable asset writes without GPU dependencies.""" + +import importlib.util +import json +from pathlib import Path +import subprocess +import sys +import tempfile +import unittest +from unittest import mock + +SCRIPT = Path(__file__).resolve().parents[2] / "recipe/jev_vl/preencode.py" +SPEC = importlib.util.spec_from_file_location("jev_vl_preencode", SCRIPT) +preencode = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(preencode) + + +class ImageManifestTests(unittest.TestCase): + def setUp(self): + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name) + self.manifest = self.root / "manifest.jsonl" + + def images(self, *states): + self.manifest.write_text("\n".join( + json.dumps({"request": {"state": state}}) for state in states + ) + "\n") + return preencode.images_of(self.manifest) + + def test_recognizes_and_deduplicates_both_image_forms(self): + first, second = "data:image/png;base64,AA==", "data:image/png;base64,AQ==" + self.assertEqual(self.images( + ["text", {"image": second}, {"type": "image_url", "image_url": {"url": first}}], + [{"image": first}], + ), [first, second]) + + def test_structured_text_does_not_become_an_image(self): + self.assertEqual(self.images( + [{"image_url": "https://example.org/logo.png"}, + {"image_url": {"url": "https://example.org/logo.png"}}, + {"type": "text", "text": "Ready.", "image_url": {"url": "ignored"}}, + {"type": "metadata", "image_url": {"url": "ignored"}}], + {"image": "top-level dictionaries are text"}, + "plain text", None, + ), []) + + def test_image_shorthand_takes_precedence(self): + self.assertEqual(self.images([{ + "type": "image_url", "image": "selected", "image_url": {"url": "ignored"}, + }]), ["selected"]) + for value in (None, "", False, 1, [], {}): + with self.subTest(value=value), self.assertRaisesRegex(ValueError, "image part"): + self.images([{"image": value, "type": "image_url", "image_url": {"url": "ignored"}}]) + + def test_rejects_malformed_typed_image_parts(self): + for image_url in (None, "bad", {}, {"url": None}, {"url": 1}, {"url": ""}): + with self.subTest(image_url=image_url), self.assertRaisesRegex(ValueError, "image part"): + self.images([{"type": "image_url", "image_url": image_url}]) + + def test_help_and_text_only_manifest_need_no_model_dependencies(self): + self.images([{"image_url": "ordinary metadata"}]) + commands = [ + ["--help"], + ["--model", str(self.root / "missing-model"), "--manifest", str(self.manifest), + "--out", str(self.root / "assets")], + ] + for args in commands: + with self.subTest(args=args): + result = subprocess.run( + [sys.executable, "-S", str(SCRIPT), *args], + capture_output=True, text=True, timeout=10, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("0 unique images", result.stdout) + self.assertFalse((self.root / "assets").exists()) + + +class AssetWriteTests(unittest.TestCase): + def setUp(self): + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name) + self.url = "data:image/png;base64,AA==" + self.key = preencode.hashlib.sha256(self.url.encode()).hexdigest() + self.dest = self.root / "assets" / self.key + self.manifest = self.root / "manifest.jsonl" + self.manifest.write_text(json.dumps({"request": {"state": [{"image": self.url}]}})) + (self.root / "model.safetensors.index.json").write_text("{}") + + tensor = mock.MagicMock(shape=(1, 5120)) + tensor.detach.return_value.cpu.return_value.to.return_value = tensor + tensor.contiguous.return_value = tensor + torch = mock.MagicMock(bfloat16="bfloat16") + torch.prod.return_value.item.return_value = 4 + torch.isfinite.return_value.all.return_value = True + model = mock.MagicMock(spatial_merge_size=2) + model.return_value.pooler_output = tensor + processor = mock.MagicMock() + processor.return_value = { + "image_grid_thw": [mock.Mock(tolist=lambda: [1, 2, 2])], + "pixel_values": mock.Mock(), + } + transformers = mock.Mock() + transformers.AutoProcessor.from_pretrained.return_value = processor + self.save = mock.Mock(side_effect=lambda tensors, filename: Path(filename).write_bytes(b"rows")) + for patcher in ( + mock.patch.dict(sys.modules, { + "torch": torch, "safetensors": mock.Mock(), + "safetensors.torch": mock.Mock(save_file=self.save), "transformers": transformers, + }), + mock.patch.object(preencode, "load_vision", return_value=model), + mock.patch.object(preencode, "decode_url", return_value=object()), + mock.patch.object(sys, "argv", [ + str(SCRIPT), "--model", str(self.root), "--manifest", str(self.manifest), + "--out", str(self.root / "assets"), + ]), + ): + patcher.start() + self.addCleanup(patcher.stop) + + def assert_complete(self): + self.assertEqual((self.dest / "emb.safetensors").read_bytes(), b"rows") + grid = json.loads((self.dest / "grid.json").read_text()) + self.assertEqual(grid["url_sha256"], self.key) + self.assertEqual(grid["grid_thw"], [1, 2, 2]) + self.assertEqual(grid["shape"], [1, 5120]) + + def test_rebuilds_incomplete_assets(self): + for existing in ((), ("emb.safetensors",), ("grid.json",)): + with self.subTest(existing=existing): + self.dest.mkdir(parents=True, exist_ok=True) + for child in self.dest.iterdir(): + child.unlink() + for name in existing: + (self.dest / name).write_bytes(b"incomplete") + self.save.reset_mock() + preencode.main() + self.save.assert_called_once() + self.assert_complete() + + def test_retries_interrupted_metadata_write(self): + original_write = Path.write_text + + def interrupted_write(path, text, *args, **kwargs): + if path.parent == self.dest: + original_write(path, text[:1], *args, **kwargs) + raise OSError("interrupted metadata write") + return original_write(path, text, *args, **kwargs) + + with mock.patch.object(Path, "write_text", interrupted_write): + with self.assertRaisesRegex(OSError, "interrupted metadata write"): + preencode.main() + self.assertFalse((self.dest / "grid.json").exists()) + preencode.main() + self.assertEqual(self.save.call_count, 2) + self.assert_complete() + + def test_retries_failed_embedding_write_with_stale_grid(self): + self.dest.mkdir(parents=True) + (self.dest / "grid.json").write_text("{}") + successful_save = self.save.side_effect + + def interrupted_save(tensors, filename): + Path(filename).write_bytes(b"partial rows") + raise OSError("interrupted embedding write") + + self.save.side_effect = interrupted_save + with self.assertRaisesRegex(OSError, "interrupted embedding write"): + preencode.main() + self.assertFalse((self.dest / "grid.json").exists()) + self.save.side_effect = successful_save + preencode.main() + self.assert_complete() + + def test_preserves_complete_asset(self): + preencode.main() + expected = {p.name: p.read_bytes() for p in self.dest.iterdir()} + self.save.reset_mock() + preencode.main() + self.save.assert_not_called() + self.assertEqual({p.name: p.read_bytes() for p in self.dest.iterdir()}, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/benchmarks/test_jev_vl_replay.py b/tests/benchmarks/test_jev_vl_replay.py new file mode 100644 index 0000000..f00561b --- /dev/null +++ b/tests/benchmarks/test_jev_vl_replay.py @@ -0,0 +1,157 @@ +"""Exercise the portable CLI against local HTTP, without model or GPU execution.""" + +import copy +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +from pathlib import Path +import statistics +import subprocess +import sys +import tempfile +import threading +import unittest + +RUNNER = Path(__file__).resolve().parents[2] / "recipe/jev_vl/replay.py" +BODY = {"kind": "noul", "effective_kind": "noul", "options": ["false", "true"], + "probabilities": [0.2, 0.8], "choice_index": 1, "choice": "true"} + + +class ReplayCliTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.manifest = self.root / "manifest.jsonl" + self.request = {"kind": "noul", "state": "Synthetic test.", "question": "Ready?"} + self.entries = [{"id": "txt-1", "request": self.request}] + self.manifest.write_text(json.dumps(self.entries[0]) + "\n") + self.reference = self.root / "reference" + self.reference.mkdir() + self.write_reference() + self.calls = [] + self.status, self.response = 200, json.dumps(BODY) + owner = self + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + payload = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + owner.calls.append((self.path, payload)) + body = owner.response.encode() + self.send_response(owner.status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + if owner.status == 302: + self.send_header("Location", "/redirected") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=self.server.serve_forever, daemon=True) + thread.start() + self.addCleanup(self.server.server_close) + self.addCleanup(self.server.shutdown) + self.base = f"http://127.0.0.1:{self.server.server_port}" + + def write_reference(self, body=BODY, status=200): + (self.reference / "txt-1.json").write_text( + json.dumps({"id": "txt-1", "status": status, "body": body}) + ) + + def run_cli(self, name="run", *extra): + return subprocess.run( + [sys.executable, str(RUNNER), "--base", self.base, + "--manifest", str(self.manifest), "--reference", str(self.reference), + "--out", str(self.root / name), *extra], + capture_output=True, text=True, timeout=15, + ) + + def read_results(self, name="run"): + out = self.root / name + return (json.loads((out / "summary.json").read_text()), + [json.loads(line) for line in (out / "results.jsonl").read_text().splitlines()]) + + def test_replays_unchanged_requests_and_excludes_per_pass_warmup(self): + self.manifest.write_text("\n".join(json.dumps(row) for row in self.entries + [ + {"id": "img-1", "request": self.request}]) + "\n") + result = self.run_cli("run", "--pattern", "txt-", "--warmup", "3") + self.assertEqual(result.returncode, 0, result.stderr) + summary, rows = self.read_results() + self.assertEqual((summary["measured"], summary["warmup"]), (2, 6)) + self.assertEqual(self.calls, [("/v1/systemone", self.request)] * 8) + self.assertEqual(summary["mean_s"], statistics.fmean( + row["client_s"] for row in rows if row["phase"] == "measured")) + self.assertTrue(all(row["max_diff"] == 0 and row["probabilities"] == [0.2, 0.8] for row in rows)) + self.assertTrue(all(json.loads(row["raw_body"]) == BODY for row in rows)) + + def test_http_and_malformed_probabilities_fail_with_raw_evidence(self): + altered = [] + for field, value in [("probabilities", []), ("probabilities", [1.0]), + ("probabilities", [0.1, 0.9]), ("kind", "score"), + ("choice_index", 0), ("choice", "false")]: + body = copy.deepcopy(BODY) + body[field] = value + altered.append((200, json.dumps(body))) + alternate = BODY | {"options": ["a", "b", "c"], "probabilities": [0.1, 0.1, 0.8], + "choice_index": 2, "choice": "c"} + altered += [(200, json.dumps(alternate)), (503, json.dumps(BODY)), (302, "{}"), + (200, "not json"), (200, "null")] + for bad in ("NaN", "Infinity", "1e400"): + altered.append((200, json.dumps(BODY).replace("0.2", bad))) + for index, (self.status, self.response) in enumerate(altered): + with self.subTest(status=self.status, response=self.response): + result = self.run_cli(f"bad-{index}", "--warmup", "0", "--passes", "1") + self.assertNotEqual(result.returncode, 0) + summary, rows = self.read_results(f"bad-{index}") + self.assertFalse(summary["ok"]) + self.assertEqual(summary["failures"], 1) + self.assertEqual(rows[0]["raw_body"], self.response) + self.assertGreater(rows[0]["client_s"], 0) + self.assertEqual(len(self.calls), len(altered)) # A redirect is never followed. + + def test_warmup_failure_stops_before_measured_requests(self): + self.status = 500 + result = self.run_cli() + self.assertNotEqual(result.returncode, 0) + summary, rows = self.read_results() + self.assertEqual((summary["measured"], summary["warmup"]), (0, 1)) + self.assertEqual(len(self.calls), 1) + self.assertEqual(rows[0]["phase"], "warmup") + + def test_probability_drift_within_fixed_tolerance_passes(self): + self.response = json.dumps(BODY | {"probabilities": [0.22, 0.78]}) + result = self.run_cli("run", "--warmup", "0", "--passes", "1") + self.assertEqual(result.returncode, 0, result.stderr) + summary, _ = self.read_results() + self.assertAlmostEqual(summary["max_diff"], 0.02) + + def test_rejects_empty_selection_empty_input_and_bad_reference_before_http(self): + self.assertNotEqual(self.run_cli("zero-passes", "--passes", "0").returncode, 0) + self.assertNotEqual(self.run_cli("negative-warmup", "--warmup", "-1").returncode, 0) + result = self.run_cli("empty-selection", "--pattern", "missing-") + self.assertNotEqual(result.returncode, 0) + self.manifest.write_text("") + self.assertNotEqual(self.run_cli("empty-input").returncode, 0) + self.manifest.write_text(json.dumps(self.entries[0]) + "\n") + for index, body in enumerate([BODY | {"probabilities": [float("nan"), 0.8]}, + BODY | {"probabilities": [1.0]}]): + self.write_reference(body) + self.assertNotEqual(self.run_cli(f"reference-{index}").returncode, 0) + self.write_reference(status=500) + self.assertNotEqual(self.run_cli("reference-status").returncode, 0) + self.assertEqual(self.calls, []) + + def test_refuses_to_overwrite_existing_results(self): + out = self.root / "run" + out.mkdir() + marker = out / "summary.json" + marker.write_text("preserve me") + self.assertNotEqual(self.run_cli().returncode, 0) + self.assertEqual(marker.read_text(), "preserve me") + self.assertEqual(self.calls, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/jev_vl/contract.rs b/tests/jev_vl/contract.rs new file mode 100644 index 0000000..72c296c --- /dev/null +++ b/tests/jev_vl/contract.rs @@ -0,0 +1,211 @@ +//! Golden pins for the JEV-27B-VL System-1 prompt contract and the upstream +//! error envelope; the strings mirror serve_decide.py's s1_pass exactly. + +use omni_jev_vl_native::contract::{Kind, compile}; + +fn labels() -> Vec { + let all: Vec = ('A'..='Z') + .map(|c| c.to_string()) + .chain(('A'..='Z').flat_map(|a| ('A'..='Z').map(move |b| format!("{a}{b}")))) + .collect(); + all[..256].to_vec() +} + +fn compile_str(raw: &str) -> omni_jev_vl_native::contract::Compiled { + compile(raw.as_bytes(), &labels()).unwrap() +} + +#[test] +fn noul_prompt_matches_serve_decide() { + let c = compile_str( + r#"{"kind":"noul","state":"Dialog: Update installed.","question":"Close it?"}"#, + ); + assert_eq!(c.kind, Kind::Noul); + assert_eq!(c.options, ["false", "true"]); + assert_eq!( + c.prompt, + "[kind] noul\n[state] Dialog: Update installed.\n[question] Close it?\n[options]\nfalse\ntrue\n[decision]:" + ); +} + +#[test] +fn score_prompt_uses_digits_lines() { + let c = compile_str(r#"{"kind":"score","state":"x","question":"Rate it."}"#); + assert_eq!(c.options, ["0", "1", "2", "3", "4", "5"]); + assert_eq!( + c.prompt, + "[kind] score\n[state] x\n[question] Rate it.\n[options]\n0\n1\n2\n3\n4\n5\n[decision]:" + ); +} + +#[test] +fn choice_prompt_labels_options() { + let c = compile_str( + r#"{"kind":"choice","state":"Dialog: x.","question":"Pick one.","options":["yes","no","later"]}"#, + ); + assert_eq!( + c.prompt, + "[kind] choice\n[state] Dialog: x.\n[question] Pick one.\n[options]\nA) yes\nB) no\nC) later\n[decision]:" + ); +} + +#[test] +fn dict_state_renders_like_python_json_dumps() { + let c = compile_str(r#"{"kind":"noul","state":{"b":1,"a":"x"},"question":"q?"}"#); + // Python json.dumps(ensure_ascii=False): ", " and ": " separators, order kept. + assert!( + c.prompt.contains(r#"[state] {"b": 1, "a": "x"}"#), + "{}", + c.prompt + ); +} + +#[test] +fn list_state_parts_concatenate_without_separator() { + let c = compile_str( + r#"{"kind":"score","state":["one","two",{"type":"text","text":"three"}],"question":"q?"}"#, + ); + assert!(c.prompt.contains("[state] onetwothree\n"), "{}", c.prompt); +} + +#[test] +fn image_state_becomes_placeholder_and_records_url() { + let c = compile_str( + r#"{"kind":"noul","state":["shot",{"image":"https://x/y.png"}],"question":"q?"}"#, + ); + assert_eq!(c.images, ["https://x/y.png"]); + assert!( + c.prompt + .contains("[state] shot<|vision_start|><|image_pad|><|vision_end|>\n"), + "{}", + c.prompt + ); +} + +#[test] +fn image_url_part_matches_upstream_image_shorthand() { + let shorthand = compile_str( + r#"{"kind":"noul","state":["shot",{"image":"data:image/png;base64,AAAA"}],"question":"q?"}"#, + ); + let typed = compile_str( + r#"{"kind":"noul","state":["shot",{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}],"question":"q?"}"#, + ); + assert_eq!(typed.images, shorthand.images); + assert_eq!(typed.prompt, shorthand.prompt); +} + +#[test] +fn system2_only_cannot_silently_return_system1() { + for value in ["true", "1", r#""true""#] { + let r = reject(&format!( + r#"{{"kind":"noul","question":"q?","system2_only":{value}}}"#, + )); + assert_eq!(r.status, 400); + assert_eq!( + r.body["error"]["message"], + "system2_only is not implemented by this worker" + ); + } + compile_str(r#"{"kind":"noul","question":"q?","system2_only":false}"#); +} + +fn reject(raw: &str) -> omni_jev_vl_native::contract::Reject { + compile(raw.as_bytes(), &labels()).unwrap_err() +} + +#[test] +fn choice_option_count_error_keeps_upstream_message() { + let r = reject(r#"{"kind":"choice","state":"x","question":"q?","options":["only"]}"#); + assert_eq!(r.status, 400); + assert_eq!( + r.body["error"]["message"], + "choice needs 2-256 options, got 1" + ); + assert_eq!(r.body["error"]["type"], "BadRequestError"); + assert_eq!(r.body["error"]["code"], 400); +} + +#[test] +fn kind_literal_error_mirrors_pydantic_envelope() { + let r = reject(r#"{"kind":"bogus","state":"x","question":"q?"}"#); + assert_eq!(r.status, 400); + assert_eq!( + r.body["error"]["message"], + "1 validation error:\n {'type': 'literal_error', 'loc': 'body.kind', 'msg': \"Input should be 'noul', 'score' or 'choice'\", 'input': 'bogus', 'ctx': {'expected': \"'noul', 'score' or 'choice'\"}}" + ); + assert_eq!(r.body["error"]["type"], "Bad Request"); + assert_eq!(r.body["error"]["param"], "kind"); + assert_eq!(r.body["error"]["code"], 400); +} + +#[test] +fn missing_question_mirrors_pydantic_envelope() { + let r = reject(r#"{"kind":"noul","state":"x"}"#); + assert_eq!( + r.body["error"]["message"], + "1 validation error:\n {'type': 'missing', 'loc': 'body.question', 'msg': 'Field required', 'input': ''}" + ); + assert_eq!(r.body["error"]["param"], "question"); +} + +#[test] +fn score_thinking_on_keeps_upstream_message() { + let r = reject(r#"{"kind":"score","state":"x","question":"q?","thinking":"on"}"#); + assert_eq!( + r.body["error"]["message"], + "thinking is supported for noul and choice" + ); +} + +#[test] +fn decision_controls_reject_non_strings() { + for field in ["thinking", "strategy"] { + for value in ["true", "null", "42", "[]", "{}"] { + let r = reject(&format!( + r#"{{"kind":"noul","question":"q?","{field}":{value}}}"#, + )); + assert_eq!(r.status, 400, "{field}: {value}"); + } + } +} + +#[test] +fn supported_decision_controls_remain_valid() { + for thinking in ["default", "off"] { + for strategy in ["auto", "single"] { + compile_str(&format!( + r#"{{"kind":"noul","question":"q?","thinking":"{thinking}","strategy":"{strategy}"}}"#, + )); + } + } +} + +#[test] +fn wrong_content_type_body_mirrors_bytes_error() { + let raw = br#"{"kind":"noul","state":"x","question":"q?"}"#; + // Non-JSON content type: upstream validates the raw bytes object. + let r = omni_jev_vl_native::contract::Reject::validation( + format!( + "1 validation error:\n {{'type': 'model_attributes_type', 'loc': 'body', 'msg': 'Input should be a valid dictionary or object to extract fields from', 'input': ''}}", + raw.len() + ), + "body", + ); + assert_eq!( + r.body["error"]["message"], + "1 validation error:\n {'type': 'model_attributes_type', 'loc': 'body', 'msg': 'Input should be a valid dictionary or object to extract fields from', 'input': ''}" + ); +} + +#[test] +fn unknown_model_keeps_404_envelope() { + let r = reject(r#"{"model":"no-such-model","kind":"noul","question":"q?"}"#); + assert_eq!(r.status, 404); + assert_eq!( + r.body["error"]["message"], + "The model `no-such-model` does not exist." + ); + assert_eq!(r.body["error"]["type"], "NotFoundError"); + assert_eq!(r.body["error"]["param"], "model"); + assert_eq!(r.body["error"]["code"], 404); +} diff --git a/tests/jev_vl/http.rs b/tests/jev_vl/http.rs new file mode 100644 index 0000000..6ed2c38 --- /dev/null +++ b/tests/jev_vl/http.rs @@ -0,0 +1,50 @@ +//! Worker fallback contract only: the frontend does not forward this route. + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::time::Duration; + +use super::*; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unknown_model_chat_probe_returns_canonical_worker_error() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + // This route needs no model or GPU; exercise the production fallback over HTTP. + let app = Router::new().fallback(fallback); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let response = tokio::task::spawn_blocking(move || { + let body = r#"{"model":"no-such-model","messages":[{"role":"user","content":"hi"}],"max_tokens":1}"#; + let mut stream = TcpStream::connect(address).unwrap(); + stream.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + stream.set_write_timeout(Some(Duration::from_secs(5))).unwrap(); + write!( + stream, + "POST /v1/chat/completions HTTP/1.0\r\nHost: {address}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ) + .unwrap(); + let mut response = String::new(); + stream.read_to_string(&mut response).unwrap(); + response + }) + .await + .unwrap(); + server.abort(); + let (headers, body) = response.split_once("\r\n\r\n").unwrap(); + assert_eq!(headers.split_whitespace().nth(1), Some("404")); + assert!( + headers + .to_ascii_lowercase() + .contains("content-type: application/json") + ); + assert_eq!( + serde_json::from_str::(body).unwrap(), + json!({"error": { + "message": "The model `no-such-model` does not exist.", + "type": "NotFoundError", + "param": "model", + "code": 404 + }}) + ); +} diff --git a/tests/jev_vl/images.rs b/tests/jev_vl/images.rs new file mode 100644 index 0000000..da54b54 --- /dev/null +++ b/tests/jev_vl/images.rs @@ -0,0 +1,101 @@ +//! Pins for the token expansion + 3-axis positions at the image boundary, +//! hand-computed from HF modeling_qwen3_5 get_rope_index semantics +//! (meshgrid t-outer/h-mid/w-fastest + start; advance = max(grid_h, grid_w)/merge). + +use std::sync::Arc; + +use omni_jev_vl_native::images::{ImageAsset, expand, meshgrid_positions}; + +const PAD: u32 = 248056; + +fn asset(grid: [i64; 3]) -> Arc { + let [t, h, w] = grid; + let n = (t * h * w / 4) as usize; + Arc::new(ImageAsset { + grid_thw: grid, + embeddings: vec![half::bf16::ONE; n * 5120], + }) +} + +#[test] +fn single_image_positions_follow_hf() { + // text 2 tokens, image grid [1,4,6] -> n=6 (lg_h=2, lg_w=3), text 1 token. + // advance = max(4, 6) / 2 = 3 -> last text at position 2 + 3 = 5. + let text_ids: Vec = vec![300, 301, PAD, 302]; + let e = expand(&text_ids, PAD, &[asset([1, 4, 6])]).unwrap(); + assert_eq!(e.ids, [300, 301, PAD, PAD, PAD, PAD, PAD, PAD, 302]); + assert_eq!(e.positions[0], vec![0, 1, 2, 2, 2, 2, 2, 2, 5]); + assert_eq!(e.positions[1], vec![0, 1, 2, 2, 2, 3, 3, 3, 5]); + assert_eq!(e.positions[2], vec![0, 1, 2, 3, 4, 2, 3, 4, 5]); + // R2d cache-anchor geometry. + assert_eq!( + e.blocks + .iter() + .map(|b| (b.start, b.end, b.base, b.advance, b.asset)) + .collect::>(), + vec![(2, 8, 2, 3, 0)] + ); +} + +#[test] +fn two_images_continue_after_max_hw() { + // hand-computed: see module doc. + let text_ids: Vec = vec![300, PAD, 301, 302, PAD, 303]; + let e = expand(&text_ids, PAD, &[asset([1, 2, 4]), asset([1, 4, 2])]).unwrap(); + assert_eq!(e.ids, [300, PAD, PAD, 301, 302, PAD, PAD, 303]); + assert_eq!(e.positions[0], vec![0, 1, 1, 3, 4, 5, 5, 7]); + assert_eq!(e.positions[1], vec![0, 1, 1, 3, 4, 5, 6, 7]); + assert_eq!(e.positions[2], vec![0, 1, 2, 3, 4, 5, 5, 7]); + assert_eq!( + e.blocks + .iter() + .map(|b| (b.start, b.end, b.base, b.advance, b.asset)) + .collect::>(), + vec![(1, 3, 1, 2, 0), (5, 7, 5, 2, 1)] + ); +} + +#[test] +fn meshgrid_helper_matches_expand_rows() { + // The helper must produce exactly expand's rows for every suffix of one block. + let text_ids: Vec = vec![300, 301, PAD, 302]; + let e = expand(&text_ids, PAD, &[asset([1, 4, 6])]).unwrap(); + let b = &e.blocks[0]; + for from in 0..6usize { + let got = meshgrid_positions([1, 4, 6], b.base, from, 6 - from); + for (want, have) in e.positions.iter().zip(&got) { + assert_eq!(&want[b.start + from..b.end], &have[..], "slice from {from}"); + } + } +} + +#[test] +fn placeholder_count_must_match_images() { + let e = expand(&[PAD], PAD, &[]); + assert!(e.is_err()); +} + +#[test] +fn image_assets_reject_wrong_sources_and_invalid_grids() { + let dir = tempfile::tempdir().unwrap(); + for (grid, url, model) in [ + ([1, 2, 2], "other-url", "model"), + ([1, 2, 2], "url", "other-model"), + ([1, 3, 2], "url", "model"), + ([2, 2, 2], "url", "model"), + ([1, i64::MAX - 1, 4], "url", "model"), + ] { + std::fs::write( + dir.path().join("grid.json"), + serde_json::to_vec(&serde_json::json!({ + "grid_thw": grid, "n_tokens": 1, + "url_sha256": url, "model_index_sha256": model, + })) + .unwrap(), + ) + .unwrap(); + let error = ImageAsset::load(dir.path(), "url", "model").err().unwrap(); + // Reject metadata before opening or allocating the embedding tensor. + assert!(!format!("{error:#}").contains("imgcache asset emb")); + } +} diff --git a/tests/jev_vl/prefix.rs b/tests/jev_vl/prefix.rs new file mode 100644 index 0000000..d39cdbc --- /dev/null +++ b/tests/jev_vl/prefix.rs @@ -0,0 +1,327 @@ +//! R2d cache-structure tests: key stability, LRU/budget accounting, and the +//! manifest-faithful suffix-split equivalence (bit-exact vs full expansion on +//! the FROZEN R1 same-image prompts — the L1 hit path's correctness contract). + +use std::path::PathBuf; +use std::sync::Arc; + +use omni_jev_vl_native::caches::{CacheCfg, Caches, structure_key}; +use omni_jev_vl_native::contract::{self, Part}; +use omni_jev_vl_native::images::{ImageAsset, expand}; + +// Compile the real cache implementation against a CPU-only prefix allocation. +// PrefixState's device buffers are private, and these tests exercise cache +// ownership/accounting, not CUDA operations or prefix numerical equivalence. +extern crate self as omni_qwen3_5_native; + +pub mod model { + pub struct PrefixState { + pub bytes: usize, + } + + impl PrefixState { + pub fn bytes(&self) -> usize { + self.bytes + } + } +} + +pub mod images { + pub use omni_jev_vl_native::images::ImageAsset; +} + +#[allow(dead_code)] +#[path = "../../src/models/jev_vl/native/src/caches.rs"] +mod cache_accounting; + +fn accounting_hub(budget: usize) -> Arc { + cache_accounting::Caches::new(cache_accounting::CacheCfg { + enabled: true, + l1: true, + l2: true, + l3: true, + l1_max: 64, + l2_bytes: 1 << 20, + l3_bytes: budget, + }) +} + +fn accounting_meta() -> cache_accounting::L1Meta { + cache_accounting::L1Meta { + pads_start: 10, + pads_end: 106, + p: 64, + base_pad: 10, + advance: 12, + ids_prefix: vec![0; 64], + positions_prefix: [vec![0; 64], vec![0; 64], vec![0; 64]], + } +} + +#[test] +fn publishing_prefix_enforces_budget_immediately() { + let cache = accounting_hub(2048); + for key in 1..=3 { + cache.record_insert(key, accounting_meta()); + cache.record_publish_state(key, model::PrefixState { bytes: 1024 }); + } + assert_eq!(cache.snapshot().l3_bytes, 2048); + assert!(cache.record_get(1).is_none()); + assert!(cache.record_get(2).unwrap().state().is_some()); + assert!(cache.record_get(3).unwrap().state().is_some()); +} + +#[test] +fn oversized_prefix_is_not_retained() { + let cache = accounting_hub(1024); + let record = cache.record_insert(1, accounting_meta()); + cache.record_publish_state(1, model::PrefixState { bytes: 2048 }); + assert!(record.state().is_none()); + assert_eq!(cache.snapshot().l3_bytes, 0); +} + +#[test] +fn zero_budget_disables_state_retention() { + let cache = accounting_hub(0); + let record = cache.record_insert(1, accounting_meta()); + cache.record_publish_state(1, model::PrefixState { bytes: 1024 }); + assert!(record.state().is_none()); + assert_eq!(cache.snapshot().l3_bytes, 0); +} + +#[test] +fn zero_record_limit_retains_no_structure() { + let cache = hub(|cfg| cfg.l1_max = 0); + cache.record_insert( + 1, + omni_jev_vl_native::caches::L1Meta { + pads_start: 10, + pads_end: 106, + p: 64, + base_pad: 10, + advance: 12, + ids_prefix: vec![0; 64], + positions_prefix: [vec![0; 64], vec![0; 64], vec![0; 64]], + }, + ); + assert!(cache.record_get(1).is_none()); + assert_eq!(cache.snapshot().l1_records, 0); +} + +#[test] +fn repeated_publication_and_eviction_preserve_live_state() { + let cache = accounting_hub(1024); + let record = cache.record_insert(1, accounting_meta()); + cache.record_publish_state(1, model::PrefixState { bytes: 1024 }); + let in_flight = record.state().unwrap(); + cache.record_publish_state(1, model::PrefixState { bytes: 1024 }); + assert_eq!(cache.snapshot().l3_bytes, 1024); + assert!(Arc::ptr_eq(&in_flight, &record.state().unwrap())); + + cache.record_insert(2, accounting_meta()); + cache.record_publish_state(2, model::PrefixState { bytes: 1024 }); + assert!(cache.record_get(1).is_none()); + assert_eq!(cache.snapshot().l3_bytes, 1024); + assert_eq!(in_flight.bytes(), 1024); +} + +#[test] +fn stats_and_prefix_publication_finish_concurrently() { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Barrier, mpsc}; + use std::time::Duration; + + let cache = accounting_hub(64 * 1024); + let start = Arc::new(Barrier::new(2)); + let done = Arc::new(AtomicBool::new(false)); + let (tx, rx) = mpsc::channel(); + let publisher = cache.clone(); + let publisher_start = start.clone(); + let publisher_done = done.clone(); + let publisher_tx = tx.clone(); + std::thread::spawn(move || { + publisher_start.wait(); + for key in 0..20_000 { + publisher.record_insert(key, accounting_meta()); + publisher.record_publish_state(key, model::PrefixState { bytes: 1024 }); + } + publisher_done.store(true, Ordering::Release); + publisher_tx.send(()).unwrap(); + }); + std::thread::spawn(move || { + start.wait(); + while !done.load(Ordering::Acquire) { + cache.snapshot(); + } + tx.send(()).unwrap(); + }); + for _ in 0..2 { + rx.recv_timeout(Duration::from_secs(10)) + .expect("cache publication and stats must not deadlock"); + } +} + +fn asset(grid: [i64; 3]) -> Arc { + let [t, h, w] = grid; + let n = (t * h * w / 4) as usize; + Arc::new(ImageAsset { + grid_thw: grid, + embeddings: vec![half::bf16::ONE; n * 5120], + }) +} + +#[test] +fn structure_key_is_sensitive_and_stable() { + let parts_a = vec![Part::Text("state text".into()), Part::Image("u://a".into())]; + let parts_b = vec![Part::Text("state text".into()), Part::Image("u://b".into())]; + let parts_c = vec![Part::Image("u://a".into()), Part::Text("state text".into())]; + assert_ne!( + structure_key("noul", &parts_a), + structure_key("score", &parts_a) + ); + assert_ne!( + structure_key("noul", &parts_a), + structure_key("noul", &parts_b) + ); + assert_ne!( + structure_key("noul", &parts_a), + structure_key("noul", &parts_c) + ); + assert_eq!( + structure_key("noul", &parts_a), + structure_key("noul", &parts_a) + ); + assert_ne!( + structure_key("noul", &parts_a), + structure_key("noul", &parts_a[..1]) + ); +} + +fn hub(f: impl FnOnce(&mut CacheCfg)) -> Arc { + let mut cfg = CacheCfg { + enabled: true, + l1: true, + l2: true, + l3: true, + l1_max: 64, + l2_bytes: 1 << 20, + l3_bytes: 1 << 20, + }; + f(&mut cfg); + Caches::new(cfg) +} + +#[test] +fn l2_lru_hits_misses_and_evicition() { + let c = hub(|cfg| cfg.l2_bytes = (2 * 5120 * 2 + 64) * 2 + 128); + let (a, b, d) = (asset([1, 2, 4]), asset([1, 2, 4]), asset([1, 2, 4])); + assert!(c.l2_get("k-a").is_none()); + c.l2_insert("k-a".into(), a.clone()); + assert!(c.l2_get("k-a").is_some()); + c.l2_insert("k-b".into(), b); + c.l2_insert("k-d".into(), d); // a evicted: budget fits two + assert!(c.l2_get("k-a").is_none()); + assert!(c.l2_get("k-b").is_some()); + let s = c.snapshot(); + assert_eq!(s.l2_hit, 2); + assert_eq!(s.l2_miss, 2); + assert_eq!(s.l2_records, 2); + // Disabled: no reuse, all misses counted. + let off = hub(|cfg| cfg.enabled = false); + off.l2_insert("k-a".into(), asset([1, 2, 4])); + assert!(off.l2_get("k-a").is_none()); + assert_eq!(off.snapshot().l2_miss, 1); +} + +#[test] +fn records_are_lru_and_budgeted() { + let c = hub(|cfg| cfg.l1_max = 2); + let meta = |pads: usize| omni_jev_vl_native::caches::L1Meta { + pads_start: pads, + pads_end: pads + 96, + p: 64, + base_pad: pads as i64, + advance: 12, + ids_prefix: vec![0u32; 64], + positions_prefix: [vec![0i64; 64], vec![0i64; 64], vec![0i64; 64]], + }; + let r1 = c.record_insert(1, meta(10)); + c.record_insert(2, meta(20)); + assert!(c.record_get(1).is_some()); + assert!(Arc::ptr_eq(&c.record_get(1).unwrap(), &r1)); + c.record_insert(3, meta(30)); // l1_max=2: one eviction of the oldest (2) + assert!(c.record_get(2).is_none()); + assert!(c.record_get(1).is_some()); + let s = c.snapshot(); + assert!(s.l1_hit >= 3 && s.l1_miss >= 1 && s.l1_records <= 2); +} + +/// The manifest-faithful equivalence: for every img entry of the FROZEN R1 +/// manifest the suffix split (pads-tail + vision_end + fresh tail tokenization) +/// must equal `expand` on the whole prompt bit for bit — that's the L1 hit path. +#[test] +#[ignore = "needs JEV_VL_EXPORT (export dir with tokenizer + manifest)"] +fn manifest_suffix_split_matches_full_expand() { + let dir = PathBuf::from(std::env::var_os("JEV_VL_EXPORT").expect("set JEV_VL_EXPORT")); + let manifest_dir = + PathBuf::from(std::env::var_os("JEV_VL_MANIFEST").expect("set JEV_VL_MANIFEST")); + let tokenizer = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).unwrap(); + let manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(dir.join("jev_vl_export.json")).unwrap()).unwrap(); + let labels: Vec = serde_json::from_value(manifest["labels"].clone()).unwrap(); + let image_pad = tokenizer.token_to_id("<|image_pad|>").unwrap(); + let vision_end = tokenizer.token_to_id("<|vision_end|>").unwrap(); + let mut checked = 0; + for line in std::fs::read_to_string(manifest_dir).unwrap().lines() { + let entry: serde_json::Value = serde_json::from_str(line).unwrap(); + if !entry["id"].as_str().unwrap_or_default().starts_with("img-") { + continue; + } + let raw = serde_json::to_vec(&entry["request"]).unwrap(); + let compiled = contract::compile(&raw, &labels).unwrap(); + let tail = contract::tail_after_last_image(&compiled, &labels).unwrap(); + let grid = [1, 60, 60]; + let e = expand( + tokenizer + .encode(compiled.prompt.as_str(), false) + .unwrap() + .get_ids(), + image_pad, + &[asset(grid)], + ) + .unwrap(); + assert_eq!(e.blocks.len(), 1, "{}", entry["id"]); + let p = e.blocks[0].end / 64 * 64; + // Rebuild as the L1 hit path does. + let b = &e.blocks[0]; + let suffix_pads = b.end - p; + let tail_ids: Vec = tokenizer + .encode(tail.as_str(), false) + .unwrap() + .get_ids() + .to_vec(); + let mut ids: Vec = std::iter::repeat_n(image_pad, suffix_pads).collect(); + ids.push(vision_end); + ids.extend_from_slice(&tail_ids); + assert_eq!(&ids[..], &e.ids[p..], "suffix ids diverge: {}", entry["id"]); + let pos = + omni_jev_vl_native::images::meshgrid_positions(grid, b.base, p - b.start, suffix_pads); + for (a, axis) in pos.iter().enumerate() { + assert_eq!( + &axis[..], + &e.positions[a][p..p + suffix_pads], + "suffix meshgrid diverges: {}", + entry["id"] + ); + } + assert_eq!( + e.positions[0][p + suffix_pads], + b.base + b.advance, + "after-image base: {}", + entry["id"] + ); + // The structure key is shared across the 12 same-image entries per kind. + checked += 1; + } + assert_eq!(checked, 12, "the frozen manifest carries 12 img entries"); +} diff --git a/tests/jev_vl/processing.rs b/tests/jev_vl/processing.rs new file mode 100644 index 0000000..c245350 --- /dev/null +++ b/tests/jev_vl/processing.rs @@ -0,0 +1,95 @@ +use super::*; +use crate::caches::CacheCfg; +use tokenizers::{ + AddedToken, models::wordlevel::WordLevel, pre_tokenizers::whitespace::WhitespaceSplit, +}; + +fn processor() -> Processor { + let vocab = [("[UNK]".to_owned(), 0)].into_iter().collect(); + let mut tokenizer = Tokenizer::new( + WordLevel::builder() + .vocab(vocab) + .unk_token("[UNK]".into()) + .build() + .unwrap(), + ); + tokenizer.with_pre_tokenizer(Some(WhitespaceSplit)); + tokenizer.add_special_tokens(&[ + AddedToken::from("<|vision_start|>", true), + AddedToken::from("<|image_pad|>", true), + AddedToken::from("<|vision_end|>", true), + ]); + let caches = Caches::new(CacheCfg { + enabled: true, + l1: true, + l2: true, + l3: false, + l1_max: 8, + l2_bytes: 1 << 20, + l3_bytes: 0, + }); + caches.l2_insert( + Processor::url_key("prepared://image"), + Arc::new(ImageAsset { + grid_thw: [1, 16, 16], + embeddings: vec![half::bf16::ONE; 64 * 5120], + }), + ); + Processor { + image_pad: tokenizer.token_to_id("<|image_pad|>").unwrap(), + vision_end: tokenizer.token_to_id("<|vision_end|>").unwrap(), + tokenizer, + labels: vec!["A".into(), "B".into()], + max_length: 256, + imgcache: None, + model_index_hash: "test".into(), + caches, + } +} + +#[test] +fn image_placeholders_are_rejected_before_and_after_l1_warmup() { + let processor = processor(); + let valid = serde_json::json!({ + "kind": "choice", "state": [{"image": "prepared://image"}], + "question": "Pick one.", "options": ["yes", "no"], + }); + let prepare = |request: &Value| { + let compiled = + contract::compile(&serde_json::to_vec(request).unwrap(), &processor.labels).unwrap(); + processor.prepare_cached( + &compiled, + Readout { + token_ids: vec![0, 1], + bias: vec![0.0; 2], + temperature: 1.0, + }, + Instant::now(), + ) + }; + let mut bad_question = valid.clone(); + bad_question["question"] = serde_json::json!("Pick <|image_pad|>."); + let mut bad_option = valid.clone(); + bad_option["options"][0] = serde_json::json!("<|image_pad|>"); + for request in [&bad_question, &bad_option] { + assert_eq!( + prepare(request) + .err() + .expect("cold request must fail") + .status, + 400 + ); + } + assert_eq!(processor.caches.snapshot().l1_records, 0); + assert_eq!(prepare(&valid).unwrap().cache_note, "l1=miss,l3=off,p=-"); + assert_eq!(prepare(&valid).unwrap().cache_note, "l1=hit,l3=off,p=-"); + for request in [&bad_question, &bad_option] { + assert_eq!( + prepare(request) + .err() + .expect("warm request must fail") + .status, + 400 + ); + } +} diff --git a/tests/jev_vl/readout.rs b/tests/jev_vl/readout.rs new file mode 100644 index 0000000..9437b33 --- /dev/null +++ b/tests/jev_vl/readout.rs @@ -0,0 +1,66 @@ +//! Frozen numeric pins for the verbalizer readout math (constants computed with +//! CPython over the same f64 formula against this test's pattern). + +use super::{LabelHead, Readout}; + +fn head() -> LabelHead { + // rows[3][64]; row t element i = ((t*64+i)*5 % 11 - 5) / 7. + let mut rows = Vec::with_capacity(3 * 64); + for t in 0..3u32 { + for i in 0..64u32 { + rows.push((((t * 64 + i) * 5 % 11) as i32 - 5) as f32 / 7.0); + } + } + LabelHead { + width: 64, + rows, + index: [(15, 0), (16, 1), (17, 2)].into_iter().collect(), + } +} + +#[test] +fn probabilities_match_python_f64() { + let hidden: Vec = (0..64).map(|i| ((i * 7 % 13) - 6) as f32 / 9.0).collect(); + let readout = Readout { + token_ids: vec![15, 16, 17], + bias: vec![0.001, -0.002, 0.0005], + temperature: 1.0143134751376188, + }; + let p = head().probabilities(&hidden, &readout).unwrap(); + let expected = [0.6256388379121548, 0.1869499071806715, 0.18741125490717364]; + // 1e-12: f32 division rounding order differs between CPython and Rust patterns; + // both far below the f32 head precision, so this pins the formula, not ulps. + for (&a, &e) in p.iter().zip(&expected) { + assert!((a - e).abs() < 1e-12, "{a} != {e}"); + } +} + +#[test] +fn softmax_is_shift_invariant_to_minus_logz() { + // The official readout subtracts per-request -logZ before bias/T; the stable + // softmax must be invariant, which is why logits suffice. + let hidden: Vec = (0..64).map(|i| ((i * 7 % 13) - 6) as f32 / 9.0).collect(); + let readout = |shift: f64| Readout { + token_ids: vec![15, 16, 17], + bias: vec![0.001 + shift, -0.002 + shift, 0.0005 + shift], + temperature: 1.0143134751376188, + }; + let p0 = head().probabilities(&hidden, &readout(0.0)).unwrap(); + let p1 = head().probabilities(&hidden, &readout(-42.5)).unwrap(); + for (&a, &b) in p0.iter().zip(&p1) { + assert!((a - b).abs() < 1e-12, "{a} != {b}"); + } +} + +#[test] +fn missing_exported_label_is_rejected() { + let mut h = head(); + h.index.remove(&17); + let readout = Readout { + token_ids: vec![15, 16, 17], + bias: vec![0.0; 3], + temperature: 1.0, + }; + let hidden = vec![0.0f32; 64]; + assert!(h.probabilities(&hidden, &readout).is_err()); +} diff --git a/tests/qwen3_5/kernels.rs b/tests/qwen3_5/kernels.rs index 99c9b8e..7a217c8 100644 --- a/tests/qwen3_5/kernels.rs +++ b/tests/qwen3_5/kernels.rs @@ -4,6 +4,9 @@ //! //! CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ //! cargo test --release -p omni-qwen3-5-native --test kernels -- --ignored +//! +//! The prefix-state test additionally needs QWEN3_5_CHECKPOINT pointing to a +//! supported multimodal checkpoint; it loads two instances sequentially. use std::path::PathBuf; @@ -617,3 +620,469 @@ fn gated_delta_rule_matches_recurrent_reference() { assert!(worst <= 2e-2 * scale, "t = {t}: {worst} vs scale {scale}"); } } + +fn f32_from_device(buf: &DeviceBuffer, n: usize, st: Stream) -> Vec { + let mut bytes = vec![0u8; n * 4]; + // SAFETY: the buffer holds n float32 values. + unsafe { cuda::download(&mut bytes, buf.at(0), st).unwrap() }; + let (quad, _) = bytes.as_chunks::<4>(); + quad.iter().map(|&b| f32::from_le_bytes(b)).collect() +} + +#[test] +#[ignore = "needs a GPU, CUA_S1_CUDA_LIB, and a multimodal QWEN3_5_CHECKPOINT"] +fn prefix_state_rejects_uncaptured_foreign_and_oversized_prompts() { + use omni_qwen3_5_native::inputs::MultimodalInput; + use omni_qwen3_5_native::model::Model; + + let library = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").expect("set CUA_S1_CUDA_LIB")); + let checkpoint = + PathBuf::from(std::env::var_os("QWEN3_5_CHECKPOINT").expect("set QWEN3_5_CHECKPOINT")); + let mut first = Model::load(&checkpoint, &library).unwrap(); + for len in [(first.cfg.max_positions / 64 + 1) * 64, usize::MAX & !63] { + let error = first.alloc_prefix(len).err().expect("oversized prefix"); + assert!(error.to_string().contains("exceeds the configured maximum")); + } + let mut state = first.alloc_prefix(64).unwrap(); + let ids = vec![0u32; state.token_count()]; + let positions: Vec = (0..state.token_count() as i64).collect(); + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&positions; 3], + }; + let error = first + .forward_multimodal_continue(&input, &state) + .unwrap_err(); + assert!(error.to_string().contains("has not completed capture")); + + first + .forward_multimodal_capture(&input, &mut state) + .unwrap(); + // The suffix fits by itself; the cached prefix pushes the full prompt over + // the limit. Reject it before allocating scratch or submitting CUDA work. + let suffix_ids = vec![0u32; first.cfg.max_positions]; + let suffix_positions = vec![0i64; suffix_ids.len()]; + let suffix = MultimodalInput { + token_ids: &suffix_ids, + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&suffix_positions; 3], + }; + let error = first + .forward_multimodal_continue(&suffix, &state) + .unwrap_err(); + assert!(error.to_string().contains("cached prompt exceeds")); + + // Retain only the prefix while replacing the model; never hold two sets of + // checkpoint weights on the GPU. The old identity must remain distinct. + drop(first); + let mut second = Model::load(&checkpoint, &library).unwrap(); + let error = second + .forward_multimodal_capture(&input, &mut state) + .unwrap_err(); + assert!(error.to_string().contains("different model instance")); + let error = second + .forward_multimodal_continue(&input, &state) + .unwrap_err(); + assert!(error.to_string().contains("different model instance")); +} + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn gdn_empty_window_captures_state_copy_and_zero_initial_state() { + let st = setup(); + let h = 2usize; + let n = h * 128 * 128; + let state = f32_to_device(&vec![1.0; n], st); + let output = f32_to_device(&vec![2.0; n], st); + let empty_window = |input: *const std::ffi::c_void| { + // SAFETY: T=0 touches only the optional input and output state buffers, + // both sized [H,128,128] FP32; token and workspace pointers are unused. + unsafe { + check( + (api().cs1_gdn_prefill_x)( + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null_mut(), + 0, + h as i32, + 1, + 1.0, + input, + output.at(0), + st, + ), + "empty gdn window", + ) + } + }; + let copy = cuda::Graph::capture(st, || empty_window(state.at(0))).unwrap(); + let updated: Vec = vec![3.0f32; n] + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect(); + // Change the source after capture: a default-stream copy during capture + // cannot substitute for the copy that must execute on every graph launch. + // SAFETY: state holds exactly n float32 values. + unsafe { cuda::upload(state.at(0), &updated, st).unwrap() }; + copy.launch(st).unwrap(); + assert_eq!(f32_from_device(&output, n, st), vec![3.0; n]); + + let zero = cuda::Graph::capture(st, || empty_window(std::ptr::null())).unwrap(); + // SAFETY: output holds exactly n float32 values; poison it after capture. + unsafe { cuda::upload(output.at(0), &updated, st).unwrap() }; + zero.launch(st).unwrap(); + assert_eq!(f32_from_device(&output, n, st), vec![0.0; n]); +} + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn device_copy_helpers_roundtrip() { + let st = setup(); + // copy_dd keeps every bit (odd size exercises plain memcpy). + let a = to_device(&random(4099, 1, 2.0), st); + let b = DeviceBuffer::new(4099 * 2).unwrap(); + // SAFETY: same size, non-overlapping device allocations. + unsafe { cuda::copy_dd(b.at(0), a.at(0), 4099 * 2, st).unwrap() }; + assert_eq!(from_device(&a, 4099, st), from_device(&b, 4099, st)); + // copy2d copies width-sized windows out of pitched rows. + let (rows, ld, w) = (5usize, 300usize, 256usize); + let src = to_device(&random(rows * ld, 2, 1.0), st); + let dst = DeviceBuffer::new(rows * w * 2).unwrap(); + // SAFETY: dst holds rows of w, src rows of ld. + unsafe { + cuda::copy2d(dst.at(0), w * 2, src.at(0), ld * 2, w * 2, rows, st).unwrap(); + } + let full = from_device(&src, rows * ld, st); + let want: Vec = full + .chunks_exact(ld) + .flat_map(|r| &r[..w]) + .copied() + .collect(); + assert_eq!(from_device(&dst, rows * w, st), want); +} + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn gdn_conv_continuation_uses_cached_history() { + let st = setup(); + let prefix = 64usize; + let (key_dim, value_dim, heads) = (16 * 128usize, 48 * 128usize, 48usize); + let channels = 2 * key_dim + value_dim; + let ld = channels + value_dim + 2 * heads; + let weights = to_device(&vec![bf16::from_f32(0.25); channels * 4], st); + let conv = |source: &DeviceBuffer, start: usize, rows: usize| { + let q = DeviceBuffer::new(rows * key_dim * 2).unwrap(); + let k = DeviceBuffer::new(rows * key_dim * 2).unwrap(); + let v = DeviceBuffer::new(rows * value_dim * 2).unwrap(); + // SAFETY: source contains start + rows complete pitched rows; outputs + // contain rows of the corresponding Q/K/V widths. + unsafe { + check( + (api().cs1_gdn_conv)( + source.at(start * ld * 2), + ld as i32, + weights.at(0), + q.at(0), + k.at(0), + v.at(0), + rows as i32, + key_dim as i32, + value_dim as i32, + st, + ), + "gdn conv continuation", + ) + .unwrap(); + } + [ + from_device(&q, rows * key_dim, st), + from_device(&k, rows * key_dim, st), + from_device(&v, rows * value_dim, st), + ] + }; + for rows in [1usize, 2, 3, 4, 65] { + let total = prefix + rows; + // Positive inputs and taps make omission of any history row observable. + let projection: Vec = random(total * ld, 37, 0.5) + .iter() + .map(|x| bf16::from_f32(x.to_f32().abs() + 0.25)) + .collect(); + let source = to_device(&projection, st); + let full = conv(&source, 0, total); + let tail = DeviceBuffer::new(3 * ld * 2).unwrap(); + let window = DeviceBuffer::new((rows + 3) * ld * 2).unwrap(); + // SAFETY: capture three prefix rows, restore them ahead of the suffix, + // and copy the suffix into the remaining non-overlapping window rows. + unsafe { + cuda::copy_dd(tail.at(0), source.at((prefix - 3) * ld * 2), 3 * ld * 2, st).unwrap(); + cuda::copy_dd(window.at(0), tail.at(0), 3 * ld * 2, st).unwrap(); + cuda::copy_dd( + window.at(3 * ld * 2), + source.at(prefix * ld * 2), + rows * ld * 2, + st, + ) + .unwrap(); + } + let restored = conv(&window, 0, rows + 3); + let missing_history = conv(&window, 3, rows); + for (i, width) in [key_dim, key_dim, value_dim].into_iter().enumerate() { + let expected = &full[i][prefix * width..]; + assert_eq!( + &restored[i][3 * width..], + expected, + "cached conv history diverges for output {i}, suffix rows={rows}" + ); + assert_ne!( + &missing_history[i][..rows.min(3) * width], + &expected[..rows.min(3) * width], + "fixture must detect the old suffix-only conv call" + ); + } + } +} + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn windowed_attention_matches_full_pass_bit_for_bit() { + let st = setup(); + let (hq, hk, dh) = (24usize, 4usize, 256usize); + for (t, bases) in [ + (139usize, vec![64usize]), + (712, vec![64, 128]), + (972, vec![64, 896]), + (2048, vec![64, 1024]), + ] { + // gated path with a strided V buffer, exactly the model's shape + let ldv = hk * dh + 16; + let q = to_device(&random(t * hq * dh, 1, 2.0), st); + let k = to_device(&random(t * hk * dh, 2, 2.0), st); + let v = to_device(&random(t * ldv, 3, 1.0), st); + let gate = to_device(&random(t * hq * dh, 4, 3.0), st); + let full = DeviceBuffer::new(t * hq * dh * 2).unwrap(); + // SAFETY: every buffer has complete rows of the shapes above. + unsafe { + check( + (api().cs1_attention_gated)( + q.at(0), + k.at(0), + v.at(0), + ldv as i32, + gate.at(0), + full.at(0), + t as i32, + hq as i32, + hk as i32, + dh as i32, + 0.0625, + st, + ), + "full pass", + ) + .unwrap(); + } + let full_rows = from_device(&full, t * hq * dh, st); + for qb in &bases { + let win = DeviceBuffer::new(t * hq * dh * 2).unwrap(); + // SAFETY: the window reads the same buffers; rows < q_base stay unwritten. + unsafe { + check( + (api().cs1_attention_gated_prefix)( + q.at(0), + k.at(0), + v.at(0), + ldv as i32, + gate.at(0), + win.at(0), + t as i32, + hq as i32, + hk as i32, + dh as i32, + 0.0625, + *qb as i32, + st, + ), + "windowed pass", + ) + .unwrap(); + } + let win_rows = from_device(&win, t * hq * dh, st); + let mut bad = 0usize; + for row in *qb..t { + for e in 0..hq * dh { + if full_rows[row * hq * dh + e] != win_rows[row * hq * dh + e] { + bad += 1; + } + } + } + assert_eq!(bad, 0, "t = {t}, q_base = {qb}: {bad} values differ"); + eprintln!("windowed attention t = {t}, q_base = {qb}: rows >= q_base bitwise equal"); + } + } +} + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn gdn_prefill_two_phase_matches_one_shot_bit_for_bit() { + let st = setup(); + let d = 128usize; + for (t, h, hk, qk_amp, decay_scale) in [ + (107usize, 4usize, 2usize, 1.0f32, 1.0f32), + (936, 4, 2, 0.01, 1.0), + (3399, 3, 1, 0.01, 1.0), + (65, 4, 2, 0.01, 1e-9), + ] { + let k = random(t * hk * d, 12, qk_amp); + let q: Vec = k + .iter() + .zip(random(t * hk * d, 11, qk_amp)) + .map(|(k, n)| bf16::from_f32(0.8 * k.to_f32() + 0.2 * n.to_f32())) + .collect(); + let v = random(t * h * d, 13, 1.0); + let g: Vec = random(t * h, 14, 1.0) + .iter() + .map(|x| (x.to_f32() - 1.0) * decay_scale) + .collect(); + let beta: Vec = random(t * h, 15, 0.5) + .iter() + .map(|x| bf16::from_f32(x.to_f32() + 0.5)) + .collect(); + let (qd, kd, vd, gd, bd) = ( + to_device(&q, st), + to_device(&k, st), + to_device(&v, st), + f32_to_device(&g, st), + to_device(&beta, st), + ); + let floats = unsafe { (api().cs1_gdn_workspace_floats)(t as i32, h as i32) }; + let ws0 = DeviceBuffer::new(floats * 4).unwrap(); + let o_full = DeviceBuffer::new(t * h * d * 2).unwrap(); + // SAFETY: every buffer holds t rows of the widths above; workspace its size. + unsafe { + check( + (api().cs1_gdn_prefill)( + qd.at(0), + kd.at(0), + vd.at(0), + gd.at(0).cast(), + bd.at(0), + o_full.at(0), + ws0.at(0).cast(), + t as i32, + h as i32, + hk as i32, + (d as f32).powf(-0.5), + st, + ), + "one-shot prefill", + ) + .unwrap(); + } + // phase 1 over rows [0, p) with the float32 state dump, phase 2 continues. + let p = ((t as f64 * 0.6) as usize / 64 * 64).max(64); + let state = DeviceBuffer::new(h * d * d * 4).unwrap(); + let o_p1 = DeviceBuffer::new(p * h * d * 2).unwrap(); + let ws1 = + DeviceBuffer::new(unsafe { (api().cs1_gdn_workspace_floats)(p as i32, h as i32) } * 4) + .unwrap(); + // SAFETY: the p-row prefix of the same inputs; the state buffer is [H][K][V] f32. + unsafe { + check( + (api().cs1_gdn_prefill_x)( + qd.at(0), + kd.at(0), + vd.at(0), + gd.at(0).cast(), + bd.at(0), + o_p1.at(0), + ws1.at(0).cast(), + p as i32, + h as i32, + hk as i32, + (d as f32).powf(-0.5), + std::ptr::null(), + state.at(0).cast(), + st, + ), + "capture phase", + ) + .unwrap(); + } + let t2 = t - p; + let o_p2 = DeviceBuffer::new(t2 * h * d * 2).unwrap(); + let ws2 = + DeviceBuffer::new(unsafe { (api().cs1_gdn_workspace_floats)(t2 as i32, h as i32) } * 4) + .unwrap(); + // SAFETY: the t2-row suffix of the same inputs; the captured state feeds it. + unsafe { + check( + (api().cs1_gdn_prefill_x)( + qd.at(p * hk * d * 2), + kd.at(p * hk * d * 2), + vd.at(p * h * d * 2), + gd.at(p * h * 4).cast(), + bd.at(p * h * 2), + o_p2.at(0), + ws2.at(0).cast(), + t2 as i32, + h as i32, + hk as i32, + (d as f32).powf(-0.5), + state.at(0).cast_const(), + std::ptr::null_mut(), + st, + ), + "continuation phase", + ) + .unwrap(); + } + let full = from_device(&o_full, t * h * d, st); + let first = from_device(&o_p1, p * h * d, st); + let second = from_device(&o_p2, t2 * h * d, st); + assert_eq!(first, full[..p * h * d], "capture rows diverge at t = {t}"); + assert_eq!( + second, + full[p * h * d..], + "continuation rows diverge at t = {t}" + ); + eprintln!("gdn two-phase t = {t}, split = {p}: both phases bitwise equal to one-shot"); + // capture-only, T = 0: the state passes through untouched. + let passthru = DeviceBuffer::new(h * d * d * 4).unwrap(); + let empty = DeviceBuffer::new(4).unwrap(); + // SAFETY: T = 0 touches only the state buffers. + unsafe { + check( + (api().cs1_gdn_prefill_x)( + empty.at(0), + empty.at(0), + empty.at(0), + empty.at(0).cast(), + empty.at(0), + empty.at(0), + empty.at(0).cast(), + 0, + h as i32, + hk as i32, + 1.0, + state.at(0).cast_const(), + passthru.at(0).cast(), + st, + ), + "state passthrough", + ) + .unwrap(); + } + assert_eq!( + f32_from_device(&state, h * d * d, st), + f32_from_device(&passthru, h * d * d, st), + "T = 0 must copy the state through" + ); + } +}