From 162e3c5f6d7b85707df0cb98cd7c93d8e7214400 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Wed, 30 Sep 2026 22:29:57 +0800 Subject: [PATCH 1/9] open_jev: add native Rust/CUDA text worker Signed-off-by: Hongsheng Liu --- Cargo.lock | 28 +- Cargo.toml | 2 +- README.md | 5 +- recipe/README.md | 2 + recipe/cua_s1/native.md | 2 +- recipe/open_jev/example-request.json | 27 ++ recipe/open_jev/export_merged.py | 70 ++++ recipe/open_jev/native.md | 105 ++++++ src/backends/cuda/qwen3_5/README.md | 9 +- src/backends/cuda/qwen3_5/attention.cu | 55 +++- src/backends/cuda/qwen3_5/ops.h | 10 +- src/models/cua_s1/native/Cargo.toml | 3 +- src/models/cua_s1/native/src/lib.rs | 4 +- src/models/open_jev/README.md | 16 + src/models/open_jev/native/Cargo.toml | 14 + src/models/open_jev/native/LICENSE.open-jev | 21 ++ src/models/open_jev/native/src/contract.rs | 230 +++++++++++++ src/models/open_jev/native/src/engine.rs | 205 ++++++++++++ src/models/open_jev/native/src/lib.rs | 3 + src/models/open_jev/native/src/main.rs | 114 +++++++ src/models/open_jev/native/tests/contract.rs | 96 ++++++ .../open_jev/native/tests/data/contract.json | 309 ++++++++++++++++++ .../native/tests/data/tokenization.json | 173 ++++++++++ src/models/qwen3_5/native/Cargo.toml | 15 + .../{cua_s1 => qwen3_5}/native/src/cuda.rs | 6 +- .../{cua_s1 => qwen3_5}/native/src/json.rs | 0 src/models/qwen3_5/native/src/lib.rs | 5 + .../{cua_s1 => qwen3_5}/native/src/model.rs | 18 +- src/models/qwen3_5/native/tests/config.rs | 17 + .../qwen3_5/native/tests/data/config.json | 140 ++++++++ .../native/tests/kernels.rs | 122 ++++++- 31 files changed, 1780 insertions(+), 46 deletions(-) create mode 100644 recipe/open_jev/example-request.json create mode 100644 recipe/open_jev/export_merged.py create mode 100644 recipe/open_jev/native.md create mode 100644 src/models/open_jev/README.md create mode 100644 src/models/open_jev/native/Cargo.toml create mode 100644 src/models/open_jev/native/LICENSE.open-jev create mode 100644 src/models/open_jev/native/src/contract.rs create mode 100644 src/models/open_jev/native/src/engine.rs create mode 100644 src/models/open_jev/native/src/lib.rs create mode 100644 src/models/open_jev/native/src/main.rs create mode 100644 src/models/open_jev/native/tests/contract.rs create mode 100644 src/models/open_jev/native/tests/data/contract.json create mode 100644 src/models/open_jev/native/tests/data/tokenization.json create mode 100644 src/models/qwen3_5/native/Cargo.toml rename src/models/{cua_s1 => qwen3_5}/native/src/cuda.rs (97%) rename src/models/{cua_s1 => qwen3_5}/native/src/json.rs (100%) create mode 100644 src/models/qwen3_5/native/src/lib.rs rename src/models/{cua_s1 => qwen3_5}/native/src/model.rs (98%) create mode 100644 src/models/qwen3_5/native/tests/config.rs create mode 100644 src/models/qwen3_5/native/tests/data/config.json rename src/models/{cua_s1 => qwen3_5}/native/tests/kernels.rs (70%) diff --git a/Cargo.lock b/Cargo.lock index 16864a1..7cc8ffb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -901,10 +901,9 @@ dependencies = [ "anyhow", "axum", "half", - "libloading", "memmap2", + "omni-qwen3-5-native", "safetensors", - "serde", "serde_json", "tokenizers", "tokio", @@ -920,6 +919,31 @@ dependencies = [ "tokio", ] +[[package]] +name = "omni-open-jev-native" +version = "0.1.0" +dependencies = [ + "anyhow", + "axum", + "omni-qwen3-5-native", + "serde_json", + "tokenizers", + "tokio", +] + +[[package]] +name = "omni-qwen3-5-native" +version = "0.1.0" +dependencies = [ + "anyhow", + "half", + "libloading", + "memmap2", + "safetensors", + "serde", + "serde_json", +] + [[package]] name = "once_cell" version = "1.21.4" diff --git a/Cargo.toml b/Cargo.toml index 8c100b4..72a6f67 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend", "src/models/cua_s1/native"] +members = ["src/frontend", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native"] resolver = "3" diff --git a/README.md b/README.md index b99d91c..46be845 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ A community-maintained inference engine for prefill-only System1-Omni models, designed around a Rust frontend, model-owned execution, and high-performance CUDA and Metal backends. -The Rust frontend forwards requests to a separately running model worker. The Cua-S1 4B 0.2 `text` adapter has a native worker with CUDA kernels in this repository; other in-repository model engines and GPU backends are not implemented yet. +The Rust frontend forwards requests to a separately running model worker. The Cua-S1 4B 0.2 `text` adapter and Open-Jev-27B-v1.1 have native workers using shared CUDA kernels in this repository. ## Run the frontend @@ -47,7 +47,7 @@ Implementation code lives under `src/`; recipes and documentation stay at the re | [`recipe/`](recipe/) | Model setup instructions, launch commands, configuration examples, and example requests. | | [`docs/`](docs/) | Project documentation and architecture assets. | -The frontend and the Cua-S1 native worker are Cargo workspace members. The other model and backend directories currently document planned work; they do not prescribe process boundaries. +The frontend, both native workers, and their shared Qwen3.5/3.8 prefill implementation are Cargo workspace members. The other model and backend directories currently document planned work; they do not prescribe process boundaries. ## Supported models @@ -57,6 +57,7 @@ LAYA can run as an external Python worker for text requests; its in-repository m | --- | --- | | LAYA | [External worker](recipe/laya/README.md); model engine planned | | 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 | +| Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](recipe/open_jev/native.md); eager independent text candidates, full-checkpoint validation pending | CUDA and Metal coverage will be documented per model as implementations are added and validated. diff --git a/recipe/README.md b/recipe/README.md index 3795fd9..cc05180 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -6,6 +6,8 @@ the worker and connect the Rust frontend. - [Cua-S1 4B 0.2 native text worker](cua_s1/native.md): build the CUDA library and 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. Recipes contain setup, launch commands and examples. Reusable implementation code belongs under `src/`. diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index cf9e996..fed34f4 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -30,5 +30,5 @@ The request tests need no GPU; the kernel tests compare attention and the chunke ```sh cargo test -p omni-cua-s1-native CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ - cargo test --release -p omni-cua-s1-native --test kernels -- --ignored + cargo test --release -p omni-qwen3-5-native --test kernels -- --ignored ``` diff --git a/recipe/open_jev/example-request.json b/recipe/open_jev/example-request.json new file mode 100644 index 0000000..b1bef5f --- /dev/null +++ b/recipe/open_jev/example-request.json @@ -0,0 +1,27 @@ +{ + "state": "A customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.", + "questions": { + "route": { + "type": "choice", + "instructions": "Which team should handle this issue?", + "criteria": { + "billing": "Problems with charges, invoices, refunds or payments.", + "security": "Unauthorized access or account compromise.", + "technical": "Service unavailable or a software malfunction." + } + }, + "refund_review": { + "type": "noul", + "instructions": "Does this message describe a duplicate charge?" + }, + "urgency": { + "type": "score", + "instructions": "Assess the urgency using only the given evidence.", + "criteria": [ + "Routine: no service disruption or active security compromise is reported.", + "Urgent: an ongoing service disruption is reported.", + "Critical: active unauthorized access is reported." + ] + } + } +} diff --git a/recipe/open_jev/export_merged.py b/recipe/open_jev/export_merged.py new file mode 100644 index 0000000..d64cb64 --- /dev/null +++ b/recipe/open_jev/export_merged.py @@ -0,0 +1,70 @@ +"""Export the pinned Open-Jev-27B-v1.1 text backbone and scalar head on CPU. + +Use the reference environment documented in native.md. This preparation step +needs about 110 GB of host RAM and 52 GB of output storage, without a GPU. +""" + +import argparse +import json +from pathlib import Path + +import torch +from peft import PeftModel +from transformers import AutoModelForImageTextToText, AutoTokenizer + +BASE_REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" +CHECKPOINT_REVISION = "28cf73067d5b337860bbef3c85b8b82ba8730956" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--base", required=True, type=Path) + parser.add_argument("--checkpoint", required=True, type=Path) + parser.add_argument("--out", required=True, type=Path) + parser.add_argument("--max-length", type=int, default=4096) + args = parser.parse_args() + config = json.loads((args.checkpoint / "model.json").read_text()) + if config["model_id"] != "Qwen/Qwen3.8-27B" or config["revision"] != BASE_REVISION: + raise ValueError("expected Open-Jev-27B-v1.1's pinned base") + if args.out.exists(): + raise ValueError("output already exists; choose a new export directory") + if not 1 <= args.max_length <= 16384: + raise ValueError("max length must be within 1..=16384") + temperature = json.loads((args.checkpoint / "temperature.json").read_text())["temperature"] + tokenizer = AutoTokenizer.from_pretrained(args.base, local_files_only=True) + marker = "\x00OMNI_OPEN_JEV\x00" + chat = tokenizer.apply_chat_template( + [{"role": "user", "content": marker}], tokenize=False, + add_generation_prompt=True, enable_thinking=False, + ) + if chat.count(marker) != 1: + raise ValueError("expected a single-user text chat template") + prefix, suffix = chat.split(marker) + head = torch.load(args.checkpoint / "head.pt", map_location="cpu", weights_only=True) + if head["weight"].shape != (1, 5120) or head["bias"].shape != (1,): + raise ValueError("expected a 5120-wide trained scalar head") + if not all(torch.isfinite(v).all() for v in head.values()): + raise ValueError("non-finite scalar head") + full = AutoModelForImageTextToText.from_pretrained( + args.base, torch_dtype=torch.bfloat16, device_map={"": "cpu"}, + attn_implementation="sdpa", local_files_only=True, + ) + backbone = full.model.language_model + del full + backbone = PeftModel.from_pretrained(backbone, args.checkpoint / "adapter") + backbone = backbone.merge_and_unload(safe_merge=True) + backbone.save_pretrained(args.out, max_shard_size="5GB") + tokenizer.save_pretrained(args.out) + # Written last: the native worker refuses incomplete exports or plain base weights. + (args.out / "open_jev_export.json").write_text(json.dumps({ + "format": "open-jev-text-merged/1", + "model_id": config["model_id"], "base_revision": BASE_REVISION, + "checkpoint_revision": CHECKPOINT_REVISION, "temperature": temperature, + "max_length": args.max_length, "chat_prefix": prefix, "chat_suffix": suffix, + "head_weight": head["weight"].float().reshape(-1).tolist(), + "head_bias": head["bias"].float().item(), + }, allow_nan=False) + "\n") + + +if __name__ == "__main__": + main() diff --git a/recipe/open_jev/native.md b/recipe/open_jev/native.md new file mode 100644 index 0000000..0742b02 --- /dev/null +++ b/recipe/open_jev/native.md @@ -0,0 +1,105 @@ +# Open-Jev-27B-v1.1 native text worker + +The worker owns request compilation, tokenization, candidate scoring and typed +responses in Rust. It uses the native CUDA prefill implementation introduced in +[PR #19](https://github.com/ThinkFlowLab/system1-omni/pull/19), shared with Cua-S1 +under [`src/models/qwen3_5/native/`](../../src/models/qwen3_5/native/). +Python is required only to prepare the merged checkpoint. + +It supports `choice` (1–255 candidates), `score` (2–10 levels), and `noul` +(yes/no). Each candidate has an independent prompt; the last hidden state goes +through Open-Jev's trained FP32 scalar head. Noul uses logits `[0, score]`. +The saved calibration temperature is applied before normalizing each complete +question. There is no autoregressive generation. Structured state and descriptions +use Open-Jev's sorted JSON rendering; question and candidate order is preserved. + +## Prepare the checkpoint + +Use the reference dependencies from +[Open-Jev @ 3308a15](https://github.com/Zefan-Cai/Open-Jev/tree/3308a15ccd7eea1df7a37d6ddc39b023b801ba16): +PyTorch 2.8 or newer, Transformers 5.10.2, PEFT 0.19.1, Accelerate 1.13.0, +and safetensors. An optional `kernels` installation must be compatible with that +Transformers release. Run these commands from the repository root: + +```sh +hf download Qwen/Qwen3.8-27B \ + --revision 1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0 \ + --local-dir weights/Qwen3.8-27B +hf download ZefanCai/Open-Jev-27B-v1.1 \ + --revision 28cf73067d5b337860bbef3c85b8b82ba8730956 \ + --include 'package/checkpoint/*' --local-dir weights/Open-Jev-27B-v1.1 +CUDA_VISIBLE_DEVICES='' python recipe/open_jev/export_merged.py \ + --base weights/Qwen3.8-27B \ + --checkpoint weights/Open-Jev-27B-v1.1/package/checkpoint \ + --out weights/open-jev-27b-merged +``` + +CPU export needs roughly 110 GB of RAM and 52 GB of output storage. It merges +LoRA in BF16 and saves the trained head, temperature, and single-user chat +template in `open_jev_export.json`. The worker refuses a plain base checkpoint +or an incomplete export. The saved limit defaults to 4096 tokens per candidate; +`--max-length` may raise it to 16384. Oversize prompts fail before inference. + +## Build and serve + +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 both Rust workers must be rebuilt together because +the gated-attention entry point updates the library ABI to version 3. + +```sh +src/backends/cuda/qwen3_5/build.sh target/release 89 +cargo build --release --locked -p omni-open-jev-native -p omni-jev +OPEN_JEV_MODEL=weights/open-jev-27b-merged \ + target/release/omni-open-jev-native +``` + +`OPEN_JEV_HOST` and `OPEN_JEV_PORT` default to `127.0.0.1` and `8000`. +`OPEN_JEV_CUDA_LIB` overrides the default library next to the executable. +The worker loads all text weights onto visible CUDA device 0, performs a real +warmup inference, then exposes `/health` and `/v1/systemone`. +Use a reservation before any GPU command on hosts with a GPU scheduler. + +In another terminal, start the existing Rust frontend: + +```sh +OMNI_JEV_BIND=127.0.0.1:8080 OMNI_JEV_BACKEND_URL=http://127.0.0.1:8000 \ + target/release/omni-jev +curl http://127.0.0.1:8080/v1/systemone \ + -H 'Content-Type: application/json' --data-binary @recipe/open_jev/example-request.json +``` + +The worker accepts the model's base name `Qwen/Qwen3.8-27B`, `open-jev`, +`jev-latest`, and `open-jev-27b-v1.1`; the response model is the base name, +matching Open-Jev. Error wording and metadata differ from the reference service. +Requests are bounded to 4 MiB, 4096 questions, and 65536 candidate sequences. + +## Validation and optimization scope + +```sh +cargo test --workspace --locked +# Inside a GPU reservation, after building the library: +CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + cargo test --release --locked -p omni-qwen3-5-native --test kernels -- --ignored +``` + +CPU golden fixtures come from Open-Jev's request compiler and response formatter +at the revision above. Kernel tests compare attention and Gated DeltaNet with +float64 references and require exact BF16 equality between fused attention gating +and a separate gate pass. The shared kernel tests retain PR #19's +`CUA_S1_CUDA_LIB` environment variable. + +The worker reuses PR #19's fused norm, activation, QK/RoPE and chunked Gated +DeltaNet operations. Attention's sigmoid gate is fused into its output epilogue, +preserving both BF16 rounding points and removing one launch and one output +read/write pass per full-attention layer (16 layers for this model). + +This initial implementation runs one eager forward pass per candidate. Prefix +sharing, CUDA Graphs, GEMM autotuning, quantization and multimodal inference are +not implemented. No latency improvement over the reference is claimed. Merged +weights, CUDA attention/Gated DeltaNet, and the CPU head's accumulation order +can change probabilities; full-checkpoint comparisons are required before +treating this worker as an accuracy-validated replacement. + +[OpenJev-Fast](https://yiqilyu.me/open-jev-fast/) motivates the fusion approach; +its B300 measurements do not apply to this worker. diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index b84f50b..1ebb27a 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -1,9 +1,14 @@ # Qwen3.5 prefill operations -CUDA kernels for a prefill-only Qwen3.5 forward pass, built into `libqwen3_5_cuda.so` with a C interface ([`ops.h`](ops.h)), so that a Rust model engine loads it at run time and builds without a CUDA toolkit. The Cua-S1 native worker ([`src/models/cua_s1/native/`](../../../models/cua_s1/native/)) uses it and keeps the layer loop and buffers. +CUDA kernels for a prefill-only Qwen3.5/3.8 forward pass, built into `libqwen3_5_cuda.so` with a C interface ([`ops.h`](ops.h)), so that a Rust model engine loads it at run time and builds without a CUDA toolkit. The Cua-S1 and Open-Jev native workers share the layer loop and buffers in [`src/models/qwen3_5/native/`](../../../models/qwen3_5/native/). ```sh src/backends/cuda/qwen3_5/build.sh [compute capability, default 89] ``` -The norm, elementwise and q/k preparation kernels round to bfloat16 where Transformers (`modeling_qwen3_5.py`) does. Attention (FlashAttention-2 style, on tensor cores) and the chunked gated delta rule keep some intermediate results in bfloat16, as FlashAttention and flash-linear-attention do. GEMMs go through cuBLASLt with its first heuristic choice. Tensor-core kernels need sm_80 or newer; only sm_89 has been run. +The norm, elementwise and q/k preparation kernels round to bfloat16 where Transformers (`modeling_qwen3_5.py`) does. Attention (FlashAttention-2 style, on tensor cores) and the chunked gated delta rule keep some intermediate results in bfloat16, as FlashAttention and flash-linear-attention do. GEMMs go through cuBLASLt with its first heuristic choice. Tensor-core kernels need sm_80 or newer; PR #19 validated the original kernels on sm_89. The new gated epilogue still needs GPU validation. + +`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 3. diff --git a/src/backends/cuda/qwen3_5/attention.cu b/src/backends/cuda/qwen3_5/attention.cu index 8c03666..20ffcee 100644 --- a/src/backends/cuda/qwen3_5/attention.cu +++ b/src/backends/cuda/qwen3_5/attention.cu @@ -78,9 +78,11 @@ 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; +template __global__ void __launch_bounds__(THREADS) flash_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv, - bf16* __restrict__ out, int T, int Hq, int Hk, float scale_log2) { + const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk, + float scale_log2) { extern __shared__ __align__(16) unsigned char smem[]; bf16* qs = reinterpret_cast(smem); bf16* ks = qs + BM * LDS; @@ -214,11 +216,39 @@ __global__ void __launch_bounds__(THREADS) if (row >= T) continue; bf16* dst = out + ((size_t)row * Hq + h) * D + 2 * t; #pragma unroll - for (int n = 0; n < D / 8; n++) - *reinterpret_cast(dst + n * 8) = pack_bf16(o[n][2 * r] * inv[r], o[n][2 * r + 1] * inv[r]); + for (int n = 0; n < D / 8; n++) { + float a = o[n][2 * r] * inv[r], b = o[n][2 * r + 1] * inv[r]; + if constexpr (Gated) { + const size_t idx = ((size_t)row * Hq + h) * D + 2 * t + n * 8; + // Match attention -> BF16 store -> BF16 sigmoid -> multiply. + // Rounding before the multiply is required even without that store. + a = round_bf16(a) * round_bf16(sigmoid(f32(gate[idx]))); + b = round_bf16(b) * round_bf16(sigmoid(f32(gate[idx + 1]))); + } + *reinterpret_cast(dst + n * 8) = pack_bf16(a, b); + } } } +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) { + 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 (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); + return cudaGetLastError(); +} + } // namespace flash } // namespace @@ -244,17 +274,10 @@ 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) { - if (Dh != flash::D || Hk <= 0 || Hq % Hk != 0 || ldv % 8 != 0 || ldv < Hk * Dh || T < 0) - return cudaErrorInvalidValue; - if (T == 0) return cudaSuccess; - // once per process (for the device current at the first call) - static const cudaError_t configured = cudaFuncSetAttribute( - flash::flash_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, flash::SMEM_BYTES); - if (configured != cudaSuccess) return configured; - constexpr float LOG2E = 1.4426950408889634f; - flash::flash_kernel<<(stream)>>>( - static_cast(q), static_cast(k), static_cast(v), ldv, - static_cast(out), T, Hq, Hk, scale * LOG2E); - return cudaGetLastError(); + return flash::launch(q, k, v, ldv, nullptr, out, T, Hq, Hk, Dh, scale, 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); } diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 6f63488..a01a564 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -12,8 +12,8 @@ #include #include -// Bumped whenever a signature below changes. -#define CS1_ABI_VERSION 2 +// Bumped whenever the required interface below changes. +#define CS1_ABI_VERSION 3 #ifdef __cplusplus extern "C" { @@ -77,6 +77,12 @@ int cs1_attn_prep(const void* qg, const void* kr, int ld, const void* qw, const 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); +// Same attention, with the sigmoid gate fused into its epilogue. gate [T, Hq, Dh]. +// Both the attention result and sigmoid are rounded to bfloat16 before multiplying, +// exactly as cs1_attention followed by cs1_sigmoid_gate. gate must be non-null if T > 0. +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); + // x = x * sigmoid(gate), n elements. int cs1_sigmoid_gate(void* x, const void* gate, size_t n, void* stream); diff --git a/src/models/cua_s1/native/Cargo.toml b/src/models/cua_s1/native/Cargo.toml index 2e8caf6..9da20a3 100644 --- a/src/models/cua_s1/native/Cargo.toml +++ b/src/models/cua_s1/native/Cargo.toml @@ -14,10 +14,9 @@ anyhow = "1.0.100" axum = "0.8.8" half = "2.7.1" # the CUDA kernels live in libqwen3_5_cuda.so, loaded at run time -libloading = "0.8" +omni-qwen3-5-native = { path = "../../qwen3_5/native" } memmap2 = "0.9.9" safetensors = "0.8.0" -serde = "1" serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order"] } # the onig regex backend, as in the Python tokenizers wheel tokenizers = { version = "=0.22.2", default-features = false, features = ["onig"] } diff --git a/src/models/cua_s1/native/src/lib.rs b/src/models/cua_s1/native/src/lib.rs index 0d3aa5f..8332012 100644 --- a/src/models/cua_s1/native/src/lib.rs +++ b/src/models/cua_s1/native/src/lib.rs @@ -3,7 +3,5 @@ //! kernels of `src/backends/cuda/qwen3_5`, loaded at run time. pub mod contract; -pub mod cuda; +pub use omni_qwen3_5_native::{cuda, json, model}; pub mod engine; -pub mod json; -pub mod model; diff --git a/src/models/open_jev/README.md b/src/models/open_jev/README.md new file mode 100644 index 0000000..0d38f43 --- /dev/null +++ b/src/models/open_jev/README.md @@ -0,0 +1,16 @@ +# Open-Jev-27B-v1.1 + +The [native Rust/CUDA worker](../../../recipe/open_jev/native.md) uses the +pinned Qwen3.8-27B backbone, merged LoRA adapter, trained scalar decision head, +and saved calibration temperature. It supports choice, ordinal score, and yes/no +text decisions through the existing Rust frontend. + +The request compiler and response formulas follow +[Open-Jev @ 3308a15](https://github.com/Zefan-Cai/Open-Jev/tree/3308a15ccd7eea1df7a37d6ddc39b023b801ba16). +The CUDA prefill implementation is shared with Cua-S1 under +[`../qwen3_5/native/`](../qwen3_5/native/). See the recipe for preparation, +numerical limitations, validation and optimization scope. + +The Rust contract is adapted from Open-Jev's MIT-licensed code; its copyright +and license are retained in [`native/LICENSE.open-jev`](native/LICENSE.open-jev). +No model weights are distributed here. diff --git a/src/models/open_jev/native/Cargo.toml b/src/models/open_jev/native/Cargo.toml new file mode 100644 index 0000000..d980ec6 --- /dev/null +++ b/src/models/open_jev/native/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "omni-open-jev-native" +version = "0.1.0" +edition = "2024" +publish = false +description = "Native Rust/CUDA worker for Open-Jev-27B-v1.1" + +[dependencies] +anyhow = "1.0.100" +axum = "0.8.8" +omni-qwen3-5-native = { path = "../../qwen3_5/native" } +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"] } diff --git a/src/models/open_jev/native/LICENSE.open-jev b/src/models/open_jev/native/LICENSE.open-jev new file mode 100644 index 0000000..d124c8c --- /dev/null +++ b/src/models/open_jev/native/LICENSE.open-jev @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Open-Jev contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/models/open_jev/native/src/contract.rs b/src/models/open_jev/native/src/contract.rs new file mode 100644 index 0000000..d65a7a1 --- /dev/null +++ b/src/models/open_jev/native/src/contract.rs @@ -0,0 +1,230 @@ +//! Open-Jev's typed request and response contract, without generated text. +//! Adapted from Open-Jev @ 3308a15; see ../LICENSE.open-jev. + +use anyhow::{Context, Result, bail, ensure}; +use omni_qwen3_5_native::json; +use serde_json::{Map, Value, json}; + +pub const MODEL_ID: &str = "Qwen/Qwen3.8-27B"; + +#[derive(Debug, PartialEq)] +pub enum Kind { + Choice, + Score, + Noul, +} + +pub struct Question { + pub id: String, + pub kind: Kind, + pub keys: Vec, + pub prompts: Vec, + pub legend: Option, +} + +fn sorted(value: &Value) -> Value { + match value { + Value::Object(map) => { + let mut entries: Vec<_> = map.iter().collect(); + entries.sort_by(|a, b| a.0.cmp(b.0)); + Value::Object( + entries + .into_iter() + .map(|(k, v)| (k.clone(), sorted(v))) + .collect(), + ) + } + Value::Array(items) => Value::Array(items.iter().map(sorted).collect()), + _ => value.clone(), + } +} + +/// Python json.dumps(ensure_ascii=False, sort_keys=True), or the string itself. +fn render(value: &Value) -> String { + value + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| json::dumps(&sorted(value))) +} + +fn description(value: &Value) -> Result { + ensure!( + value.is_string() || value.is_object() || value.is_array(), + "instructions and descriptions must be text, an object, or an array" + ); + Ok(render(value)) +} + +pub fn compile(raw: &[u8]) -> Result> { + let request = json::parse(raw).map_err(anyhow::Error::msg)?; + ensure!( + matches!(request.get("model"), None | Some(Value::Null)) + || matches!( + request["model"].as_str(), + Some(MODEL_ID | "open-jev" | "jev-latest" | "open-jev-27b-v1.1") + ), + "requested model is not loaded" + ); + let state = request.get("state").context("request requires state")?; + ensure!( + state.is_string() || state.is_object() || state.is_array(), + "state must be text, a JSON object, or an array" + ); + let questions = request + .get("questions") + .and_then(Value::as_object) + .context("questions must be an object")?; + ensure!( + !questions.is_empty() && questions.len() <= 4096, + "questions must be nonempty and within the server question limit" + ); + let mut result = Vec::with_capacity(questions.len()); + let mut candidates = 0; + for (id, definition) in questions { + let d = definition + .as_object() + .context("question definition must be an object")?; + let mut instructions = description(d.get("instructions").unwrap_or(&Value::Null))?; + let criteria = d.get("criteria").unwrap_or(&Value::Null); + let (kind, keys, options, legend) = match d.get("type").and_then(Value::as_str) { + Some("choice") => { + let criteria = criteria.as_object().context("Choice requires an object")?; + ensure!( + (1..=255).contains(&criteria.len()), + "Choice requires between 1 and 255 candidates" + ); + let mut options = Vec::with_capacity(criteria.len()); + for (name, value) in criteria { + options.push(if value.is_null() { + name.clone() + } else { + format!("{name}: {}", description(value)?) + }); + } + ( + Kind::Choice, + criteria.keys().cloned().collect::>(), + options, + None, + ) + } + Some("score") => { + let criteria = criteria.as_array().context("Score requires an array")?; + ensure!( + (2..=10).contains(&criteria.len()), + "Score requires 2 to 10 descriptive levels" + ); + let keys: Vec = (0..criteria.len()).map(|i| i.to_string()).collect(); + let legend: Map = + keys.iter().cloned().zip(criteria.iter().cloned()).collect(); + ( + Kind::Score, + keys, + criteria + .iter() + .map(description) + .collect::>>()?, + Some(Value::Object(legend)), + ) + } + Some("noul") => { + if !criteria.is_null() { + let criteria = criteria + .as_object() + .context("Noul criteria must contain true and false descriptions")?; + ensure!( + criteria.len() == 2 + && criteria.contains_key("true") + && criteria.contains_key("false"), + "Noul criteria must contain true and false descriptions" + ); + instructions += &format!( + "\nYes means: {}\nNo means: {}", + description(&criteria["true"])?, + description(&criteria["false"])? + ); + } + ( + Kind::Noul, + vec!["false".into(), "true".into()], + Vec::new(), + None, + ) + } + _ => bail!("question type must be choice, score, or noul"), + }; + let prefix = format!("Context:\n{}\n\nQuestion: {instructions}\n", render(state)); + let prompts = if kind == Kind::Noul { + vec![prefix + "Is the answer to this question yes? Answer Yes or No."] + } else { + options.into_iter().map(|option| format!("{prefix}Proposed answer: {option}\nIs this proposed answer correct? Answer Yes or No.")).collect() + }; + candidates += prompts.len(); + ensure!(candidates <= 65536, "request exceeds the candidate limit"); + result.push(Question { + id: id.clone(), + kind, + keys, + prompts, + legend, + }); + } + Ok(result) +} + +/// Calibrated normalization is across the complete question's candidates. +pub fn answer(question: &Question, logits: &[f32], temperature: f64) -> Result { + ensure!( + temperature.is_finite() && temperature > 0.0, + "invalid temperature" + ); + ensure!( + logits.len() == question.keys.len() && logits.iter().all(|v| v.is_finite()), + "invalid logits" + ); + let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max) as f64; + let mut probs: Vec = logits + .iter() + .map(|&v| ((v as f64 - max) / temperature).exp()) + .collect(); + let total: f64 = probs.iter().sum(); + probs.iter_mut().for_each(|p| *p /= total); + if question.kind == Kind::Noul { + return Ok(json!({"type": "noul", "noul": probs[1]})); + } + // First maximum wins ties, like Python's max(..., key=...). + let mode = (1..probs.len()).fold(0, |m, i| if probs[i] > probs[m] { i } else { m }); + let probabilities: Map = question + .keys + .iter() + .cloned() + .zip(probs.iter().map(|p| json!(p))) + .collect(); + if question.kind == Kind::Choice { + let n = probs.len() as f64; + let confidence = if probs.len() == 1 { + 1.0 + } else { + (probs[mode] - 1.0 / n) / (1.0 - 1.0 / n) + }; + Ok( + json!({"type": "choice", "choice": question.keys[mode], "probabilities": probabilities, "confidence": confidence}), + ) + } else { + let score: f64 = probs.iter().enumerate().map(|(i, p)| i as f64 * p).sum(); + let distance: f64 = probs + .iter() + .enumerate() + .map(|(i, p)| i.abs_diff(mode) as f64 * p) + .sum(); + let center = (probs.len() - 1) as f64 / 2.0; + let uniform_distance = (0..probs.len()) + .map(|i| (i as f64 - center).abs()) + .sum::() + / probs.len() as f64; + Ok( + json!({"type": "score", "score": score, "probabilities": probabilities, + "confidence": (1.0 - distance / uniform_distance).max(0.0), "legend": question.legend}), + ) + } +} diff --git a/src/models/open_jev/native/src/engine.rs b/src/models/open_jev/native/src/engine.rs new file mode 100644 index 0000000..107a6dc --- /dev/null +++ b/src/models/open_jev/native/src/engine.rs @@ -0,0 +1,205 @@ +//! Independent candidates through the shared Qwen3.5/3.8 CUDA prefill. + +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use anyhow::{Context, Result, ensure}; +use omni_qwen3_5_native::model::{Config, Model}; +use serde_json::Value; +use tokenizers::Tokenizer; + +use crate::contract::{Kind, MODEL_ID, Question}; + +pub const BASE_REVISION: &str = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0"; +pub const CHECKPOINT_REVISION: &str = "28cf73067d5b337860bbef3c85b8b82ba8730956"; + +pub struct Engine { + tokenizer: Tokenizer, + model: Arc>, + head: Vec, + bias: f32, + prefix: String, + suffix: String, + pub temperature: f64, + pub max_length: usize, +} + +impl Engine { + pub async fn load(dir: &Path, library: &Path) -> Result { + let manifest: Value = serde_json::from_slice( + &std::fs::read(dir.join("open_jev_export.json")) + .context("export the merged checkpoint; see recipe/open_jev/native.md")?, + )?; + ensure!( + manifest["format"] == "open-jev-text-merged/1" + && manifest["model_id"] == MODEL_ID + && manifest["base_revision"] == BASE_REVISION + && manifest["checkpoint_revision"] == CHECKPOINT_REVISION, + "expected a pinned Open-Jev-27B-v1.1 export" + ); + let temperature = manifest["temperature"].as_f64().context("temperature")?; + ensure!( + temperature.is_finite() && temperature > 0.0, + "invalid temperature" + ); + let max_length = manifest["max_length"].as_u64().context("max_length")? as usize; + ensure!( + (1..=16384).contains(&max_length), + "max_length must be within 1..=16384" + ); + let prefix = manifest["chat_prefix"] + .as_str() + .context("chat_prefix")? + .to_owned(); + let suffix = manifest["chat_suffix"] + .as_str() + .context("chat_suffix")? + .to_owned(); + 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 head: Vec = manifest["head_weight"] + .as_array() + .context("head_weight")? + .iter() + .map(|v| { + v.as_f64() + .map(|v| v as f32) + .context("head weight must be numeric") + }) + .collect::>()?; + let bias = manifest["head_bias"].as_f64().context("head_bias")? as f32; + ensure!( + head.len() == cfg.hidden && head.iter().all(|v| v.is_finite()) && bias.is_finite(), + "invalid scalar decision head" + ); + let tokenizer = + Tokenizer::from_file(dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + 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 { + tokenizer, + model: Arc::new(Mutex::new(model)), + head, + bias, + prefix, + suffix, + temperature, + max_length, + }) + } + + /// Validate every length before starting GPU inference; never truncate. + pub fn encode(&self, questions: &[Question]) -> Result>>> { + encode_questions( + &self.tokenizer, + &self.prefix, + &self.suffix, + self.max_length, + questions, + ) + } + + pub async fn score( + &self, + ids: Vec>>, + questions: &[Question], + ) -> Result>> { + let model = self.model.clone(); + let head = self.head.clone(); + let bias = self.bias; + let mut rows = tokio::task::spawn_blocking(move || { + let mut model = model + .lock() + .map_err(|_| anyhow::anyhow!("poisoned model"))?; + ids.iter() + .map(|candidates| { + candidates + .iter() + .map(|ids| { + let last = model.forward(ids)?; + // BF16 hidden -> trained FP32 scalar head. FP64 accumulation on the + // CPU, rounded to FP32, can differ from PyTorch's FP32 GEMM order. + let score = (head + .iter() + .zip(last) + .map(|(&w, h)| w as f64 * h as f64) + .sum::() as f32) + + bias; + ensure!(score.is_finite(), "non-finite decision score"); + Ok(score) + }) + .collect::>>() + }) + .collect::>>() + }) + .await??; + for (row, question) in rows.iter_mut().zip(questions) { + if question.kind == Kind::Noul { + row.insert(0, 0.0); // logits [false=0, true=trained score] + } + } + Ok(rows) + } +} + +fn encode_questions( + tokenizer: &Tokenizer, + prefix: &str, + suffix: &str, + max_length: usize, + questions: &[Question], +) -> Result>>> { + questions + .iter() + .map(|q| { + q.prompts.iter().map(|prompt| { + let chat = format!("{prefix}{}{suffix}", prompt.trim()); + let enc = tokenizer.encode(chat, true).map_err(anyhow::Error::msg)?; + ensure!(enc.len() <= max_length, + "question {}: {} tokens exceeds max_length={max_length}; no silent truncation", + q.id, enc.len()); + Ok(enc.get_ids().to_vec()) + }).collect() + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::contract; + + #[test] + #[ignore = "needs the pinned export at OPEN_JEV_MODEL; no GPU needed"] + fn tokenization_matches_reference_and_rejects_oversize_prompts() { + let dir = std::env::var_os("OPEN_JEV_MODEL").expect("OPEN_JEV_MODEL"); + let dir = Path::new(&dir); + let manifest: Value = + serde_json::from_slice(&std::fs::read(dir.join("open_jev_export.json")).unwrap()) + .unwrap(); + let tokenizer = Tokenizer::from_file(dir.join("tokenizer.json")).unwrap(); + let prefix = manifest["chat_prefix"].as_str().unwrap(); + let suffix = manifest["chat_suffix"].as_str().unwrap(); + let cases: Value = + serde_json::from_str(include_str!("../tests/data/tokenization.json")).unwrap(); + for case in cases.as_array().unwrap() { + let questions = + contract::compile(&serde_json::to_vec(&case["request"]).unwrap()).unwrap(); + let got = encode_questions(&tokenizer, prefix, suffix, 4096, &questions).unwrap(); + assert_eq!(serde_json::json!(got), case["ids"]); + let limit = got[0][0].len() - 1; + assert!(encode_questions(&tokenizer, prefix, suffix, limit, &questions).is_err()); + } + } +} diff --git a/src/models/open_jev/native/src/lib.rs b/src/models/open_jev/native/src/lib.rs new file mode 100644 index 0000000..a8a1cf3 --- /dev/null +++ b/src/models/open_jev/native/src/lib.rs @@ -0,0 +1,3 @@ +//! Open-Jev-27B-v1.1 request compilation, candidate scoring and typed responses. +pub mod contract; +pub mod engine; diff --git a/src/models/open_jev/native/src/main.rs b/src/models/open_jev/native/src/main.rs new file mode 100644 index 0000000..f09ed04 --- /dev/null +++ b/src/models/open_jev/native/src/main.rs @@ -0,0 +1,114 @@ +//! OPEN_JEV_MODEL= omni-open-jev-native + +use std::sync::Arc; +use std::time::Instant; + +use anyhow::{Context, Result, ensure}; +use axum::{ + Json, Router, + body::Bytes, + extract::{DefaultBodyLimit, State, rejection::BytesRejection}, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use omni_open_jev_native::{ + contract::{self, MODEL_ID}, + engine::{BASE_REVISION, Engine}, +}; +use omni_qwen3_5_native::cuda; +use serde_json::{Map, json}; + +const WARMUP: &[u8] = br#"{"state":"Dialog: Update installed.","questions":{"q":{"type":"choice","instructions":"Close it.","criteria":{"ok":"OK","wait":"Wait"}}}}"#; + +fn error(status: StatusCode, message: impl ToString) -> Response { + (status, Json(json!({"error": message.to_string()}))).into_response() +} + +async fn decide(engine: &Engine, raw: &[u8]) -> Response { + let questions = match contract::compile(raw) { + Ok(q) => q, + Err(e) => return error(StatusCode::UNPROCESSABLE_ENTITY, e), + }; + let start = Instant::now(); + let ids = match engine.encode(&questions) { + Ok(ids) => ids, + Err(e) => return error(StatusCode::UNPROCESSABLE_ENTITY, e), + }; + let input_tokens: usize = ids.iter().flatten().map(Vec::len).sum(); + let candidates: usize = ids.iter().map(Vec::len).sum(); + let result = async { + let rows = engine.score(ids, &questions).await?; + let mut answers = Map::new(); + for (q, logits) in questions.iter().zip(rows) { + answers.insert(q.id.clone(), contract::answer(q, &logits, engine.temperature)?); + } + Ok::<_, anyhow::Error>(json!({"model": MODEL_ID, "answers": answers, + "usage": {"input_tokens": input_tokens, "output_tokens": 0}, + "metadata": {"method": "native_merged_lora_decision_head", "temperature": engine.temperature, + "candidate_sequences": candidates, "inference_seconds": start.elapsed().as_secs_f64(), + "base_revision": BASE_REVISION, "max_length": engine.max_length, + "prefix_cache": {"enabled": false, "mode": "independent_candidates"}}})) + }.await; + match result { + Ok(body) => Json(body).into_response(), + Err(e) => { + eprintln!("inference failed: {e:#}"); + error(StatusCode::INTERNAL_SERVER_ERROR, "model inference failed") + } + } +} + +async fn systemone( + State(engine): State>, + headers: HeaderMap, + body: Result, +) -> Response { + let content_type = headers + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .split(';') + .next() + .unwrap_or("") + .trim(); + if !content_type.eq_ignore_ascii_case("application/json") { + return error( + StatusCode::UNSUPPORTED_MEDIA_TYPE, + "Content-Type must be application/json", + ); + } + match body { + Ok(raw) => decide(&engine, &raw).await, + Err(e) => error(e.status(), e.body_text()), + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let model = std::env::var_os("OPEN_JEV_MODEL").context("set OPEN_JEV_MODEL")?; + let library = std::env::var_os("OPEN_JEV_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("OPEN_JEV_HOST").unwrap_or_else(|_| "127.0.0.1".into()); + let port: u16 = std::env::var("OPEN_JEV_PORT") + .map_or(Ok(8000), |v| v.parse()) + .context("OPEN_JEV_PORT")?; + let app = Router::new() + .route( + "/health", + get(|| async { Json(json!({"status": "ready", "model": MODEL_ID})) }), + ) + .route("/v1/systemone", post(systemone)) + .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(()) +} diff --git a/src/models/open_jev/native/tests/contract.rs b/src/models/open_jev/native/tests/contract.rs new file mode 100644 index 0000000..4c63afa --- /dev/null +++ b/src/models/open_jev/native/tests/contract.rs @@ -0,0 +1,96 @@ +use omni_open_jev_native::contract::{answer, compile}; +use serde_json::{Value, json}; + +fn close(actual: &Value, expected: &Value) { + match (actual, expected) { + (Value::Number(a), Value::Number(b)) => { + assert!( + (a.as_f64().unwrap() - b.as_f64().unwrap()).abs() < 1e-12, + "{a} != {b}" + ); + } + (Value::Object(a), Value::Object(b)) => { + assert_eq!(a.len(), b.len()); + for (key, value) in b { + close(&a[key], value); + } + } + _ => assert_eq!(actual, expected), + } +} + +#[test] +fn prompts_and_typed_answers_match_open_jev_reference() { + // Generated with jev.api and jev.metrics @ 3308a15ccd7eea1df7a37d6ddc39b023b801ba16. + let cases: Value = serde_json::from_str(include_str!("data/contract.json")).unwrap(); + for case in cases.as_array().unwrap() { + let raw = serde_json::to_vec(&case["request"]).unwrap(); + let questions = compile(&raw).unwrap(); + let prompts: Vec<_> = questions.iter().map(|q| &q.prompts).collect(); + assert_eq!(json!(prompts), case["prompts"]); + let mut answers = serde_json::Map::new(); + for (q, row) in questions.iter().zip(case["logits"].as_array().unwrap()) { + let logits: Vec = row + .as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap() as f32) + .collect(); + answers.insert( + q.id.clone(), + answer(q, &logits, case["temperature"].as_f64().unwrap()).unwrap(), + ); + } + close(&json!({"answers": answers}), &case["response"]); + } +} + +#[test] +fn rejects_malformed_and_unsupported_requests() { + for raw in [ + r#"{"state":"x","state":"y","questions":{}}"#, + r#"{"state":NaN,"questions":{}}"#, + r#"{"state":12,"questions":{"q":{"type":"noul","instructions":"x"}}}"#, + r#"{"state":"x","questions":{}}"#, + r#"{"model":"other","state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}"#, + r#"{"state":"x","questions":{"q":{"type":"generation","instructions":"x"}}}"#, + r#"{"state":"x","questions":{"q":{"type":"choice","instructions":"x","criteria":{}}}}"#, + r#"{"state":"x","questions":{"q":{"type":"choice","instructions":"x","criteria":{"a":4}}}}"#, + r#"{"state":"x","questions":{"q":{"type":"score","instructions":"x","criteria":["one"]}}}"#, + r#"{"state":"x","questions":{"q":{"type":"noul","instructions":"x","criteria":{"true":"yes"}}}}"#, + ] { + assert!(compile(raw.as_bytes()).is_err(), "accepted {raw}"); + } + let q = + compile(br#"{"state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}"#).unwrap(); + for temperature in [0.0, -1.0, f64::NAN, f64::INFINITY] { + assert!(answer(&q[0], &[0.0, 1.0], temperature).is_err()); + } + assert!(answer(&q[0], &[1.0], 1.0).is_err()); + assert!(answer(&q[0], &[0.0, f32::NAN], 1.0).is_err()); +} + +#[test] +fn choice_and_score_limits_match_reference() { + let mut candidates = serde_json::Map::new(); + for i in 0..255 { + candidates.insert(i.to_string(), Value::Null); + } + let mut body = json!({"state": "x", "questions": {"q": { + "type": "choice", "instructions": "Pick", "criteria": candidates + }}}); + let questions = compile(&serde_json::to_vec(&body).unwrap()).unwrap(); + assert_eq!(questions[0].prompts.len(), 255); + body["questions"]["q"]["criteria"]["extra"] = Value::Null; + assert!(compile(&serde_json::to_vec(&body).unwrap()).is_err()); + body["questions"]["q"] = + json!({"type": "score", "instructions": "Rate", "criteria": vec!["level"; 10]}); + assert_eq!( + compile(&serde_json::to_vec(&body).unwrap()).unwrap()[0] + .prompts + .len(), + 10 + ); + body["questions"]["q"]["criteria"] = json!(vec!["level"; 11]); + assert!(compile(&serde_json::to_vec(&body).unwrap()).is_err()); +} diff --git a/src/models/open_jev/native/tests/data/contract.json b/src/models/open_jev/native/tests/data/contract.json new file mode 100644 index 0000000..e90a417 --- /dev/null +++ b/src/models/open_jev/native/tests/data/contract.json @@ -0,0 +1,309 @@ +[ + { + "request": { + "state": "A customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.", + "questions": { + "route": { + "type": "choice", + "instructions": "Which team should handle this issue?", + "criteria": { + "billing": "Problems with charges, invoices, refunds or payments.", + "security": "Unauthorized access or account compromise.", + "technical": "Service unavailable or a software malfunction." + } + }, + "refund_review": { + "type": "noul", + "instructions": "Does this message describe a duplicate charge?" + }, + "urgency": { + "type": "score", + "instructions": "Assess the urgency using only the given evidence.", + "criteria": [ + "Routine: no service disruption or active security compromise is reported.", + "Urgent: an ongoing service disruption is reported.", + "Critical: active unauthorized access is reported." + ] + } + } + }, + "prompts": [ + [ + "Context:\nA customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.\n\nQuestion: Which team should handle this issue?\nProposed answer: billing: Problems with charges, invoices, refunds or payments.\nIs this proposed answer correct? Answer Yes or No.", + "Context:\nA customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.\n\nQuestion: Which team should handle this issue?\nProposed answer: security: Unauthorized access or account compromise.\nIs this proposed answer correct? Answer Yes or No.", + "Context:\nA customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.\n\nQuestion: Which team should handle this issue?\nProposed answer: technical: Service unavailable or a software malfunction.\nIs this proposed answer correct? Answer Yes or No." + ], + [ + "Context:\nA customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.\n\nQuestion: Does this message describe a duplicate charge?\nIs the answer to this question yes? Answer Yes or No." + ], + [ + "Context:\nA customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.\n\nQuestion: Assess the urgency using only the given evidence.\nProposed answer: Routine: no service disruption or active security compromise is reported.\nIs this proposed answer correct? Answer Yes or No.", + "Context:\nA customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.\n\nQuestion: Assess the urgency using only the given evidence.\nProposed answer: Urgent: an ongoing service disruption is reported.\nIs this proposed answer correct? Answer Yes or No.", + "Context:\nA customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.\n\nQuestion: Assess the urgency using only the given evidence.\nProposed answer: Critical: active unauthorized access is reported.\nIs this proposed answer correct? Answer Yes or No." + ] + ], + "logits": [ + [ + -2.0, + -1.0, + 0.0 + ], + [ + -2.0, + -1.0 + ], + [ + -2.0, + -1.0, + 0.0 + ] + ], + "temperature": 2.5343690298472983, + "response": { + "answers": { + "route": { + "type": "choice", + "choice": "technical", + "probabilities": { + "billing": 0.21343437444688138, + "security": 0.3166841593736535, + "technical": 0.4698814661794651 + }, + "confidence": 0.20482219926919762 + }, + "refund_review": { + "type": "noul", + "noul": 0.5973836777434065 + }, + "urgency": { + "type": "score", + "score": 1.2564470917325836, + "probabilities": { + "0": 0.21343437444688138, + "1": 0.3166841593736535, + "2": 0.4698814661794651 + }, + "confidence": 0.0, + "legend": { + "0": "Routine: no service disruption or active security compromise is reported.", + "1": "Urgent: an ongoing service disruption is reported.", + "2": "Critical: active unauthorized access is reported." + } + } + } + } + }, + { + "request": { + "state": { + "z": [ + 1.0, + 1e+16, + -0.0, + true, + null, + "猫" + ], + "a": { + "β": "z", + "a": "first" + } + }, + "questions": { + "kept-id": { + "type": "choice", + "instructions": { + "z": 1, + "a": "choose" + }, + "criteria": { + "z-last": null, + "a-first": [ + "A", + { + "z": 2, + "a": 1 + } + ] + } + } + } + }, + "prompts": [ + [ + "Context:\n{\"a\": {\"a\": \"first\", \"β\": \"z\"}, \"z\": [1.0, 1e+16, -0.0, true, null, \"猫\"]}\n\nQuestion: {\"a\": \"choose\", \"z\": 1}\nProposed answer: z-last\nIs this proposed answer correct? Answer Yes or No.", + "Context:\n{\"a\": {\"a\": \"first\", \"β\": \"z\"}, \"z\": [1.0, 1e+16, -0.0, true, null, \"猫\"]}\n\nQuestion: {\"a\": \"choose\", \"z\": 1}\nProposed answer: a-first: [\"A\", {\"a\": 1, \"z\": 2}]\nIs this proposed answer correct? Answer Yes or No." + ] + ], + "logits": [ + [ + -2.0, + -1.0 + ] + ], + "temperature": 2.5343690298472983, + "response": { + "answers": { + "kept-id": { + "type": "choice", + "choice": "a-first", + "probabilities": { + "z-last": 0.4026163222565936, + "a-first": 0.5973836777434065 + }, + "confidence": 0.194767355486813 + } + } + } + }, + { + "request": { + "state": [ + "café", + { + "b": false, + "a": 1e-05 + } + ], + "questions": { + "score": { + "type": "score", + "instructions": "Rate.", + "criteria": [ + "low", + "middle", + "high" + ] + }, + "binary": { + "type": "noul", + "instructions": [ + "Ready?", + { + "b": 2, + "a": 1 + } + ], + "criteria": { + "true": { + "z": 1, + "a": "yes" + }, + "false": [ + "no", + "later" + ] + } + } + } + }, + "prompts": [ + [ + "Context:\n[\"café\", {\"a\": 1e-05, \"b\": false}]\n\nQuestion: Rate.\nProposed answer: low\nIs this proposed answer correct? Answer Yes or No.", + "Context:\n[\"café\", {\"a\": 1e-05, \"b\": false}]\n\nQuestion: Rate.\nProposed answer: middle\nIs this proposed answer correct? Answer Yes or No.", + "Context:\n[\"café\", {\"a\": 1e-05, \"b\": false}]\n\nQuestion: Rate.\nProposed answer: high\nIs this proposed answer correct? Answer Yes or No." + ], + [ + "Context:\n[\"café\", {\"a\": 1e-05, \"b\": false}]\n\nQuestion: [\"Ready?\", {\"a\": 1, \"b\": 2}]\nYes means: {\"a\": \"yes\", \"z\": 1}\nNo means: [\"no\", \"later\"]\nIs the answer to this question yes? Answer Yes or No." + ] + ], + "logits": [ + [ + -2.0, + -1.0, + 0.0 + ], + [ + -2.0, + -1.0 + ] + ], + "temperature": 2.5343690298472983, + "response": { + "answers": { + "score": { + "type": "score", + "score": 1.2564470917325836, + "probabilities": { + "0": 0.21343437444688138, + "1": 0.3166841593736535, + "2": 0.4698814661794651 + }, + "confidence": 0.0, + "legend": { + "0": "low", + "1": "middle", + "2": "high" + } + }, + "binary": { + "type": "noul", + "noul": 0.5973836777434065 + } + } + } + }, + { + "request": { + "model": "open-jev-27b-v1.1", + "state": "Anything", + "questions": { + "one": { + "type": "choice", + "instructions": "Pick", + "criteria": { + "only": null + } + }, + "tie": { + "type": "choice", + "instructions": "Pick", + "criteria": { + "first": "first", + "second": "second" + } + } + } + }, + "prompts": [ + [ + "Context:\nAnything\n\nQuestion: Pick\nProposed answer: only\nIs this proposed answer correct? Answer Yes or No." + ], + [ + "Context:\nAnything\n\nQuestion: Pick\nProposed answer: first: first\nIs this proposed answer correct? Answer Yes or No.", + "Context:\nAnything\n\nQuestion: Pick\nProposed answer: second: second\nIs this proposed answer correct? Answer Yes or No." + ] + ], + "logits": [ + [ + 0.0 + ], + [ + 0.0, + 0.0 + ] + ], + "temperature": 2.5343690298472983, + "response": { + "answers": { + "one": { + "type": "choice", + "choice": "only", + "probabilities": { + "only": 1.0 + }, + "confidence": 1.0 + }, + "tie": { + "type": "choice", + "choice": "first", + "probabilities": { + "first": 0.5, + "second": 0.5 + }, + "confidence": 0.0 + } + } + } + } +] diff --git a/src/models/open_jev/native/tests/data/tokenization.json b/src/models/open_jev/native/tests/data/tokenization.json new file mode 100644 index 0000000..4dc56ad --- /dev/null +++ b/src/models/open_jev/native/tests/data/tokenization.json @@ -0,0 +1,173 @@ +[ + { + "request": { + "state": "A customer says: I was charged twice for one order. The service is working normally. There is no sign of unauthorized access.", + "questions": { + "route": { + "type": "choice", + "instructions": "Which team should handle this issue?", + "criteria": { + "billing": "Problems with charges, invoices, refunds or payments.", + "security": "Unauthorized access or account compromise.", + "technical": "Service unavailable or a software malfunction." + } + }, + "refund_review": { + "type": "noul", + "instructions": "Does this message describe a duplicate charge?" + }, + "urgency": { + "type": "score", + "instructions": "Assess the urgency using only the given evidence.", + "criteria": [ + "Routine: no service disruption or active security compromise is reported.", + "Urgent: an ongoing service disruption is reported.", + "Critical: active unauthorized access is reported." + ] + } + } + }, + "ids": [ + [ + [248045, 846, 198, 1905, 25, 198, 32, 5813, 2640, 25, 353, 557, 11102, 10598, 364, 799, 1906, 13, 561, 2393, 369, 3133, 13870, 13, 2532, 369, 874, 1777, 314, 42994, 2528, 13, 271, 14162, 25, 15451, 2015, 1220, 3579, 411, 4125, 30, 198, 1941, 18988, 4087, 25, 32419, 25, 43057, 440, 9892, 11, 61893, 11, 70057, 466, 13767, 13, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 32, 5813, 2640, 25, 353, 557, 11102, 10598, 364, 799, 1906, 13, 561, 2393, 369, 3133, 13870, 13, 2532, 369, 874, 1777, 314, 42994, 2528, 13, 271, 14162, 25, 15451, 2015, 1220, 3579, 411, 4125, 30, 198, 1941, 18988, 4087, 25, 4610, 25, 60972, 2528, 466, 2605, 28425, 13, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 32, 5813, 2640, 25, 353, 557, 11102, 10598, 364, 799, 1906, 13, 561, 2393, 369, 3133, 13870, 13, 2532, 369, 874, 1777, 314, 42994, 2528, 13, 271, 14162, 25, 15451, 2015, 1220, 3579, 411, 4125, 30, 198, 1941, 18988, 4087, 25, 10597, 25, 5194, 33829, 466, 264, 3061, 68770, 13, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ], + [ + [248045, 846, 198, 1905, 25, 198, 32, 5813, 2640, 25, 353, 557, 11102, 10598, 364, 799, 1906, 13, 561, 2393, 369, 3133, 13870, 13, 2532, 369, 874, 1777, 314, 42994, 2528, 13, 271, 14162, 25, 12193, 411, 1876, 7276, 264, 21814, 6545, 30, 198, 3742, 279, 4087, 310, 411, 3296, 9542, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ], + [ + [248045, 846, 198, 1905, 25, 198, 32, 5813, 2640, 25, 353, 557, 11102, 10598, 364, 799, 1906, 13, 561, 2393, 369, 3133, 13870, 13, 2532, 369, 874, 1777, 314, 42994, 2528, 13, 271, 14162, 25, 79013, 279, 52011, 1608, 1132, 279, 2574, 5721, 13, 198, 1941, 18988, 4087, 25, 68650, 25, 874, 2393, 41696, 466, 4393, 4610, 28425, 369, 4800, 13, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 32, 5813, 2640, 25, 353, 557, 11102, 10598, 364, 799, 1906, 13, 561, 2393, 369, 3133, 13870, 13, 2532, 369, 874, 1777, 314, 42994, 2528, 13, 271, 14162, 25, 79013, 279, 52011, 1608, 1132, 279, 2574, 5721, 13, 198, 1941, 18988, 4087, 25, 82534, 305, 25, 449, 13789, 2393, 41696, 369, 4800, 13, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 32, 5813, 2640, 25, 353, 557, 11102, 10598, 364, 799, 1906, 13, 561, 2393, 369, 3133, 13870, 13, 2532, 369, 874, 1777, 314, 42994, 2528, 13, 271, 14162, 25, 79013, 279, 52011, 1608, 1132, 279, 2574, 5721, 13, 198, 1941, 18988, 4087, 25, 33513, 25, 4393, 42994, 2528, 369, 4800, 13, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ] + ] + }, + { + "request": { + "state": { + "z": [ + 1.0, + 1e+16, + -0.0, + true, + null, + "猫" + ], + "a": { + "β": "z", + "a": "first" + } + }, + "questions": { + "kept-id": { + "type": "choice", + "instructions": { + "z": 1, + "a": "choose" + }, + "criteria": { + "z-last": null, + "a-first": [ + "A", + { + "z": 2, + "a": 1 + } + ] + } + } + } + }, + "ids": [ + [ + [248045, 846, 198, 1905, 25, 198, 4754, 64, 763, 5046, 64, 763, 328, 3765, 487, 328, 49531, 763, 328, 89, 13933, 328, 89, 763, 498, 16, 13, 15, 11, 220, 16, 68, 10, 16, 21, 11, 471, 15, 13, 15, 11, 804, 11, 819, 11, 328, 98500, 89254, 271, 14162, 25, 5046, 64, 763, 328, 26184, 487, 328, 89, 763, 220, 16, 92, 198, 1941, 18988, 4087, 25, 1110, 58477, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 4754, 64, 763, 5046, 64, 763, 328, 3765, 487, 328, 49531, 763, 328, 89, 13933, 328, 89, 763, 498, 16, 13, 15, 11, 220, 16, 68, 10, 16, 21, 11, 471, 15, 13, 15, 11, 804, 11, 819, 11, 328, 98500, 89254, 271, 14162, 25, 5046, 64, 763, 328, 26184, 487, 328, 89, 763, 220, 16, 92, 198, 1941, 18988, 4087, 25, 264, 35718, 25, 4241, 32, 487, 5046, 64, 763, 220, 16, 11, 328, 89, 763, 220, 17, 24634, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ] + ] + }, + { + "request": { + "state": [ + "café", + { + "b": false, + "a": 1e-05 + } + ], + "questions": { + "score": { + "type": "score", + "instructions": "Rate.", + "criteria": [ + "low", + "middle", + "high" + ] + }, + "binary": { + "type": "noul", + "instructions": [ + "Ready?", + { + "b": 2, + "a": 1 + } + ], + "criteria": { + "true": { + "z": 1, + "a": "yes" + }, + "false": [ + "no", + "later" + ] + } + } + } + }, + "ids": [ + [ + [248045, 846, 198, 1905, 25, 198, 1143, 895, 56868, 487, 5046, 64, 763, 220, 16, 68, 12, 15, 20, 11, 328, 65, 763, 867, 24634, 271, 14162, 25, 19172, 13, 198, 1941, 18988, 4087, 25, 3238, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 1143, 895, 56868, 487, 5046, 64, 763, 220, 16, 68, 12, 15, 20, 11, 328, 65, 763, 867, 24634, 271, 14162, 25, 19172, 13, 198, 1941, 18988, 4087, 25, 5956, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 1143, 895, 56868, 487, 5046, 64, 763, 220, 16, 68, 12, 15, 20, 11, 328, 65, 763, 867, 24634, 271, 14162, 25, 19172, 13, 198, 1941, 18988, 4087, 25, 1496, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ], + [ + [248045, 846, 198, 1905, 25, 198, 1143, 895, 56868, 487, 5046, 64, 763, 220, 16, 68, 12, 15, 20, 11, 328, 65, 763, 867, 24634, 271, 14162, 25, 4241, 18622, 29993, 5046, 64, 763, 220, 16, 11, 328, 65, 763, 220, 17, 24634, 198, 9175, 3254, 25, 5046, 64, 763, 328, 9405, 487, 328, 89, 763, 220, 16, 92, 198, 2665, 3254, 25, 4241, 2083, 487, 328, 65297, 1293, 198, 3742, 279, 4087, 310, 411, 3296, 9542, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ] + ] + }, + { + "request": { + "model": "open-jev-27b-v1.1", + "state": "Anything", + "questions": { + "one": { + "type": "choice", + "instructions": "Pick", + "criteria": { + "only": null + } + }, + "tie": { + "type": "choice", + "instructions": "Pick", + "criteria": { + "first": "first", + "second": "second" + } + } + } + }, + "ids": [ + [ + [248045, 846, 198, 1905, 25, 198, 74660, 271, 14162, 25, 19123, 198, 1941, 18988, 4087, 25, 1132, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ], + [ + [248045, 846, 198, 1905, 25, 198, 74660, 271, 14162, 25, 19123, 198, 1941, 18988, 4087, 25, 1118, 25, 1118, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], + [248045, 846, 198, 1905, 25, 198, 74660, 271, 14162, 25, 19123, 198, 1941, 18988, 4087, 25, 2018, 25, 2018, 198, 3742, 411, 10660, 4087, 4252, 30, 21134, 7179, 466, 2233, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271] + ] + ] + } +] diff --git a/src/models/qwen3_5/native/Cargo.toml b/src/models/qwen3_5/native/Cargo.toml new file mode 100644 index 0000000..f3e0fee --- /dev/null +++ b/src/models/qwen3_5/native/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "omni-qwen3-5-native" +version = "0.1.0" +edition = "2024" +publish = false +description = "Shared Qwen3.5/3.8 prefill implementation for native decision workers" + +[dependencies] +anyhow = "1.0.100" +half = "2.7.1" +libloading = "0.8" +memmap2 = "0.9.9" +safetensors = "0.8.0" +serde = "1" +serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order"] } diff --git a/src/models/cua_s1/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs similarity index 97% rename from src/models/cua_s1/native/src/cuda.rs rename to src/models/qwen3_5/native/src/cuda.rs index a15fef2..21ac3b1 100644 --- a/src/models/cua_s1/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 = 2; +const ABI_VERSION: u32 = 3; pub const LIBRARY: &str = "libqwen3_5_cuda.so"; /// A `cudaStream_t`. @@ -91,6 +91,10 @@ api! { q: *const c_void, k: *const c_void, v: *const c_void, ldv: c_int, 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( + 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_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; diff --git a/src/models/cua_s1/native/src/json.rs b/src/models/qwen3_5/native/src/json.rs similarity index 100% rename from src/models/cua_s1/native/src/json.rs rename to src/models/qwen3_5/native/src/json.rs diff --git a/src/models/qwen3_5/native/src/lib.rs b/src/models/qwen3_5/native/src/lib.rs new file mode 100644 index 0000000..929435b --- /dev/null +++ b/src/models/qwen3_5/native/src/lib.rs @@ -0,0 +1,5 @@ +//! Qwen3.5/3.8 prefill, runtime-loaded CUDA operations, and request JSON helpers +//! shared by the Cua-S1 and Open-Jev native workers. +pub mod cuda; +pub mod json; +pub mod model; diff --git a/src/models/cua_s1/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs similarity index 98% rename from src/models/cua_s1/native/src/model.rs rename to src/models/qwen3_5/native/src/model.rs index 2b04077..a2dc45a 100644 --- a/src/models/cua_s1/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -1,4 +1,4 @@ -//! The Qwen3.5 text model (the language model of Qwen/Qwen3.5-4B), prefill only: one +//! The Qwen3.5/3.8 text backbone, prefill only: one //! forward pass over a prompt, returning the final-norm hidden state of the last //! position. The layer loop and buffers live here; the operations are the CUDA //! kernels in `src/backends/cuda/qwen3_5`. @@ -218,7 +218,7 @@ impl Weights { .enumerate() .flat_map(|(i, st)| st.names().into_iter().map(move |n| (i, n.to_string()))) .collect(); - let prefix = ["model.language_model.", "model."] + let prefix = ["model.language_model.", "model.", ""] .into_iter() .find(|p| { names @@ -792,11 +792,12 @@ impl Model { "attention prep", )?; check( - (cuda::api().cs1_attention)( + (cuda::api().cs1_attention_gated)( p(s.aq), p(s.ak), p(v), ld, + p(s.agate), p(s.ao), ti, hq, @@ -805,16 +806,7 @@ impl Model { (cfg.head_dim as f32).powf(-0.5), st, ), - "attention", - )?; - check( - (cuda::api().cs1_sigmoid_gate)( - p(s.ao), - p(s.agate), - t * cfg.heads * cfg.head_dim, - st, - ), - "attention gate", + "gated attention", )?; } self.gemm(s, s.ao, &fa.o, s.delta, t)?; diff --git a/src/models/qwen3_5/native/tests/config.rs b/src/models/qwen3_5/native/tests/config.rs new file mode 100644 index 0000000..a6e4540 --- /dev/null +++ b/src/models/qwen3_5/native/tests/config.rs @@ -0,0 +1,17 @@ +use omni_qwen3_5_native::model::Config; +use std::path::Path; + +#[test] +fn open_jev_27b_backbone_has_supported_dimensions() { + // Qwen/Qwen3.8-27B @ 1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0. + let cfg = Config::load(Path::new(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/data" + ))) + .unwrap(); + assert_eq!((cfg.hidden, cfg.intermediate), (5120, 17408)); + assert_eq!((cfg.heads, cfg.kv_heads, cfg.head_dim), (24, 4, 256)); + assert_eq!((cfg.lin_k_heads, cfg.lin_v_heads), (16, 48)); + assert_eq!(cfg.full_attention.len(), 64); + assert_eq!(cfg.full_attention.iter().filter(|&&x| x).count(), 16); +} diff --git a/src/models/qwen3_5/native/tests/data/config.json b/src/models/qwen3_5/native/tests/data/config.json new file mode 100644 index 0000000..706cebd --- /dev/null +++ b/src/models/qwen3_5/native/tests/data/config.json @@ -0,0 +1,140 @@ +{ + "architectures": [ + "Qwen3_5ForConditionalGeneration" + ], + "image_token_id": 248056, + "language_model_only": false, + "model_type": "qwen3_5", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attn_output_gate": true, + "bos_token_id": 248044, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 5120, + "initializer_range": 0.02, + "intermediate_size": 17408, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_value_head_dim": 128, + "mamba_ssm_dtype": "float32", + "max_position_embeddings": 262144, + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 24, + "num_hidden_layers": 64, + "num_key_value_heads": 4, + "output_gate_type": "swish", + "pad_token_id": null, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + }, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 248320 + }, + "tie_word_embeddings": false, + "transformers_version": "5.8.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 27, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen3_5", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 5120, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053 +} \ No newline at end of file diff --git a/src/models/cua_s1/native/tests/kernels.rs b/src/models/qwen3_5/native/tests/kernels.rs similarity index 70% rename from src/models/cua_s1/native/tests/kernels.rs rename to src/models/qwen3_5/native/tests/kernels.rs index cd044c8..03236d8 100644 --- a/src/models/cua_s1/native/tests/kernels.rs +++ b/src/models/qwen3_5/native/tests/kernels.rs @@ -8,7 +8,7 @@ use std::path::PathBuf; use half::bf16; -use omni_cua_s1_native::cuda::{self, DeviceBuffer, Stream, api, check}; +use omni_qwen3_5_native::cuda::{self, DeviceBuffer, Stream, api, check}; fn setup() -> Stream { let lib = std::env::var_os("CUA_S1_CUDA_LIB") @@ -59,6 +59,126 @@ fn from_device(buf: &DeviceBuffer, n: usize, st: Stream) -> Vec { .collect() } +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn fused_attention_gate_matches_separate_pass() { + let st = setup(); + let dh = 256usize; + for (hq, hk) in [(24usize, 4usize), (16, 4), (4, 4), (16, 1)] { + for t in [1usize, 63, 64, 65, 139, 712, 2048] { + let ldv = hk * dh + 16; // exercise a strided V buffer + 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 n = t * hq * dh; + let mut gates = random(n, 4, 12.0); + // Saturated sigmoid, zero, and ordinary values in every output row. + for row in gates.chunks_exact_mut(dh) { + row[0] = bf16::from_f32(-40.0); + row[1] = bf16::from_f32(40.0); + row[2] = bf16::ZERO; + } + let gate = to_device(&gates, st); + let separate = DeviceBuffer::new(n * 2).unwrap(); + let fused = DeviceBuffer::new(n * 2).unwrap(); + // SAFETY: all buffers have complete rows of the widths above, with + // disjoint output and gate allocations, and use the same stream. + unsafe { + check( + (api().cs1_attention)( + q.at(0), + k.at(0), + v.at(0), + ldv as i32, + separate.at(0), + t as i32, + hq as i32, + hk as i32, + dh as i32, + 0.0625, + st, + ), + "separate attention", + ) + .unwrap(); + check( + (api().cs1_sigmoid_gate)(separate.at(0), gate.at(0), n, st), + "separate gate", + ) + .unwrap(); + check( + (api().cs1_attention_gated)( + q.at(0), + k.at(0), + v.at(0), + ldv as i32, + gate.at(0), + fused.at(0), + t as i32, + hq as i32, + hk as i32, + dh as i32, + 0.0625, + st, + ), + "fused attention gate", + ) + .unwrap(); + } + let want = from_device(&separate, n, st); + let got = from_device(&fused, n, st); + assert!( + got.iter().all(|x| x.is_finite()), + "non-finite result at t={t}" + ); + assert_eq!( + got.iter().map(|v| v.to_bits()).collect::>(), + want.iter().map(|v| v.to_bits()).collect::>(), + "fused gate changed BF16 output at t={t}, hq={hq}, hk={hk}" + ); + } + } + // Empty work is valid without a gate buffer; nonempty work must have one. + // SAFETY: these invalid/empty shapes return before any kernel is launched. + unsafe { + let null = std::ptr::null(); + assert_eq!( + (api().cs1_attention_gated)( + null, + null, + null, + 256, + null, + std::ptr::null_mut(), + 0, + 1, + 1, + 256, + 0.0625, + st + ), + 0 + ); + assert_ne!( + (api().cs1_attention_gated)( + null, + null, + null, + 256, + null, + std::ptr::null_mut(), + 1, + 1, + 1, + 256, + 0.0625, + st + ), + 0 + ); + } +} + #[test] #[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] fn flash_attention_matches_float64_reference() { From 61b83b380baca912c0dbd989d175cb3da8a881ae Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Fri, 2 Oct 2026 12:37:34 +0800 Subject: [PATCH 2/9] cuda: cache residual norm and pack MLP SiLU Signed-off-by: Hongsheng Liu --- README.md | 2 +- recipe/open_jev/native.md | 13 +- recipe/open_jev/validation.md | 87 ++++++++++++ src/backends/cuda/qwen3_5/README.md | 2 +- src/backends/cuda/qwen3_5/elementwise.cu | 24 +++- src/backends/cuda/qwen3_5/norm.cu | 37 +++++ src/models/qwen3_5/native/tests/kernels.rs | 156 ++++++++++++++++++++- 7 files changed, 311 insertions(+), 10 deletions(-) create mode 100644 recipe/open_jev/validation.md diff --git a/README.md b/README.md index 301d462..8679ec8 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,7 @@ LAYA can run as an external Python worker for text requests; its in-repository m | --- | --- | | LAYA | [External worker](recipe/laya/README.md); [CPU checkpoint reader](src/models/laya/README.md); model execution planned | | 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 | -| Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](recipe/open_jev/native.md); eager independent text candidates, full-checkpoint validation pending | +| Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](recipe/open_jev/native.md); eager independent text candidates; [L20X validation](recipe/open_jev/validation.md) | CUDA and Metal coverage will be documented per model as implementations are added and validated. diff --git a/recipe/open_jev/native.md b/recipe/open_jev/native.md index a0538c7..cf89b6e 100644 --- a/recipe/open_jev/native.md +++ b/recipe/open_jev/native.md @@ -87,19 +87,24 @@ CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ CPU golden fixtures come from Open-Jev's request compiler and response formatter at the revision above. Kernel tests compare attention and Gated DeltaNet with float64 references and require exact BF16 equality between fused attention gating -and a separate gate pass. The shared kernel tests retain PR #19's +and a separate gate pass. Residual RMSNorm and packed SiLU are checked against +rounded references, including odd widths and unaligned pointers. The shared +kernel tests retain PR #19's `CUA_S1_CUDA_LIB` environment variable. The worker reuses PR #19's fused norm, activation, QK/RoPE and chunked Gated DeltaNet operations. Attention's sigmoid gate is fused into its output epilogue, preserving both BF16 rounding points and removing one launch and one output -read/write pass per full-attention layer (16 layers for this model). +read/write pass per full-attention layer (16 layers for this model). Residual +RMSNorm keeps thread values in registers at widths 2560/5120. MLP SiLU uses +16-byte BF16 loads/stores when width, stride and pointers permit it, retaining +both BF16 rounding points; other layouts use the scalar path. See the +[L20X validation](validation.md) for the measured scope and variability. This recipe leaves `CUA_S1_GRAPH` unset and runs one eager forward pass per candidate. The shared backend retains Cua-S1's opt-in CUDA Graph path, but Open-Jev graph replay remains unvalidated. Prefix sharing, GEMM autotuning, -quantization and multimodal inference are not implemented. No latency -improvement over the reference is claimed. Merged +quantization and multimodal inference are not implemented. The L20X comparison does not establish a general speedup over OpenJev-Fast. Merged weights, CUDA attention/Gated DeltaNet, and the CPU head's accumulation order can change probabilities; full-checkpoint comparisons are required before treating this worker as an accuracy-validated replacement. diff --git a/recipe/open_jev/validation.md b/recipe/open_jev/validation.md new file mode 100644 index 0000000..8c57ffe --- /dev/null +++ b/recipe/open_jev/validation.md @@ -0,0 +1,87 @@ +# Open-Jev L20X validation, 2026-10-02 + +A matched comparison of 74 real JevBench `noul` requests, each with one candidate, +validated the shared CUDA residual RMSNorm and packed MLP SiLU changes. +All 74 native decisions and probabilities were exactly unchanged by packed SiLU; +the native worker scored 64/74 correct. OpenJev-Fast scored 63/74, with one different +decision (`hard-opus-a-temporal_numeric-09`) and maximum native/Fast probability +difference 0.03435. These counts do not establish statistical accuracy superiority. + +## Warm HTTP latency + +| Configuration | Mean, pass 1 / pass 2 (ms) | P50, pass 1 / pass 2 (ms) | +| --- | --- | --- | +| Native with cached residual RMSNorm; scalar SiLU | 49.563 /49.490 | 25.434 /25.539 | +| Native with cached residual RMSNorm and packed SiLU | 48.286 /48.307 | 24.925 /24.657 | +| Original OpenJev-Fast | 57.345 /50.908 | 25.026 /24.770 | + +Packed SiLU reduces native mean 49.527→48.297 ms (2.483%). Its two-pass mean range +is disjoint from baseline's. These are observed two-pass ranges, not confidence +intervals. Fast's 6.436 ms mean spread prevents claiming a stable aggregate winner. +This subset does not measure multi-candidate prefix sharing. The author's 17.3 ms +B300 result is a different hardware/workload measurement. + +HTTP includes the same Rust frontend, tokenization and worker execution; excludes +client body serialization and response JSON parsing, includes UTF8 decoding. One +server per configuration is reused. Real warmup, first request after readiness and +one complete 74-case feasibility pass are excluded; two subsequent passes are +measured. Nsight collection is inactive during HTTP passes, although CUPTI +instrumentation may remain loaded. No shared caches are dropped or clocks changed. + +## Separate CUDA timelines + +Each entry is the two-trace mean of 64 MLP SiLU launches per request, in milliseconds. + +| Tokens | Scalar SiLU | Packed SiLU | Fast SiLU lookup | +| --- | --- | --- | --- | +| 107 | 0.680 | 0.305 | 0.359 | +| 936 | 5.643 | 1.746 | 1.513 | +| 3399 | 21.217 | 6.380 | 5.466 | + +Long-request MLP SiLU decreases 69.93%, closing 94.19% of that measured kernel +family's gap to Fast. Other kernel families vary between traces; their separate +timeline totals are not HTTP latency. Fast uses padded/tree layouts and lookup +tables, so this does not imply identical rows, fusion boundaries or arithmetic. + +The prespecified acceptance gates were five GPU tests, unchanged native decisions, +maximum probability delta ≤0.01, ≥25% long SiLU duration reduction and ≥2% warm HTTP +mean reduction with disjoint observed ranges. All passed; no extra measured runs +were added. An earlier TF32 inverse candidate failed the probability gate; a +512-thread GDN candidate missed the performance gates. Neither is included here. + +## Frozen controls and reproduction + +- Device: exact scheduler GPU 2, UUID `GPU-cbf66259-f4ab-0ede-1811-82037dde5924`, + NVIDIA L20X 143771 MiB; CUDA driver and Nsight identify SM90 / 132 SMs. NUMA 0, CPUs 0–15. +- BF16, max length 16384, HTTP concurrency 1, native `CUA_S1_GRAPH=0`. Fast retains + its original graph/kernel stack. Build with nvcc 13.0.88, SM90, `-O3 -std=c++17 + -lineinfo`; use the same cuBLASLt/runtime libraries for both native variants. +- Frozen native worker/frontend: `202c0e163f868334a99d88407056ebe61dbb2dce`. + Both native libraries include the cached RMSNorm source; only `elementwise.cu` + differs for packed SiLU. The shared CUDA ABI remains 4. +- [JevBench](https://github.com/fstandhartinger/jevbench) + at `f8ce71361165846101d02ebc83ad44e47ae44fc3`; select its 74 `noul` cases, preserving + request bodies and order. Frozen request JSON SHA256: + `0c756a7b0b4c1f1352225f2e01770b5b3a0646fe86d5cc6930ded9e292ef37df`. +- [OpenJev-Fast](https://github.com/lyuyiqi/open-jev-fast) + at `c52b8bb958c1f0d241d4eb7fce4ecd8d885bf1e4`; original server/model/kernels, with + prepared SM90 extensions, PyTorch 2.13cu130/Triton 3.7.1, Transformers 5.10.2, + PEFT 0.19.1 and FLA 0.5.2. These differ from the author's B300 environment. +- Model/export uses the pinned base, adapter and calibration in the + [native recipe](native.md); temperature 2.5343690298472983. Reuse prepared weights + and extensions. Keep copies, downloads, compilation, warmup and process-to- + readiness separate from measured execution. + +Reserve the exact device with `gpu run --gpu-ids 2 --timeout 45m --note