diff --git a/Cargo.lock b/Cargo.lock
index 4b9adda..715fbd3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -949,10 +949,9 @@ dependencies = [
"anyhow",
"axum",
"half",
- "libloading",
"memmap2",
+ "omni-qwen3-5-native",
"safetensors 0.8.0",
- "serde",
"serde_json",
"tokenizers",
"tokio",
@@ -982,6 +981,31 @@ dependencies = [
"tempfile",
]
+[[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 0.8.0",
+ "serde",
+ "serde_json",
+]
+
[[package]]
name = "once_cell"
version = "1.21.4"
diff --git a/Cargo.toml b/Cargo.toml
index 661c83e..e7f3326 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,3 +1,3 @@
[workspace]
-members = ["src/frontend", "src/models/cua_s1/native", "src/models/laya"]
+members = ["src/frontend", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/laya"]
resolver = "3"
diff --git a/README.md b/README.md
index dfe39f9..cbf99e6 100644
--- a/README.md
+++ b/README.md
@@ -27,9 +27,17 @@ 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.
+Cua-S1 4B 0.2 `text` adapter and Open-Jev-27B-v1.1 have native workers using
+shared CUDA kernels in this repository.
+
+## News
+
+- **2026-10-03:** Added [Open-Jev-27B-v1.1](recipe/open_jev/native.md)
+ support through a native Rust/CUDA worker: **7.47× faster than raw HF Transformers**
+ by mean warm HTTP latency, **362.21→48.50 ms** on one H200. Measured over
+ 74 single-candidate JevBench `noul` requests per pass, with two measured passes
+ per backend (BF16, concurrency 1). See the
+ [HF Transformers baseline, results and OpenJev-Fast comparison](recipe/open_jev/validation.md).
## Features
@@ -38,9 +46,9 @@ implemented yet.
model workers.
- **Model-owned execution.** Each model owns its preprocessing, batching,
state, execution, and kernel selection; shared utilities stay minimal.
-- **Native CUDA worker.** The Cua-S1 4B 0.2 `text` adapter runs as a native
- worker with CUDA kernels, or as a Python worker that serves as the
- correctness reference.
+- **Native CUDA workers.** The Cua-S1 4B 0.2 `text` adapter and
+ Open-Jev-27B-v1.1 run as native workers with shared CUDA kernels. Cua-S1
+ also has a Python worker that serves as the correctness reference.
- **LAYA text serving.** LAYA runs as an external Python worker for text
requests, with an in-repository CPU checkpoint reader.
- **CUDA and Metal backends.** High-performance GPU operations for NVIDIA
@@ -80,9 +88,10 @@ repository root.
| [`recipe/`](recipe/) | Model setup instructions, launch commands, configuration examples, and example requests. |
| [`docs/`](docs/) | Project documentation and architecture assets. |
-The frontend, Cua-S1 native worker and Laya checkpoint reader 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, their shared Qwen3.5/3.8 prefill
+implementation and the Laya checkpoint reader are Cargo workspace members.
+The other model and backend directories currently document planned work;
+they do not prescribe process boundaries.
## Getting Started
@@ -104,12 +113,14 @@ for a CPU text worker and response checks, or the Cua-S1 recipes for the
LAYA can run as an external Python worker for text requests; its
in-repository model engine is still planned. The Cua-S1 4B 0.2 `text` adapter
-runs as a Python worker or as a native worker on CUDA:
+runs as a Python worker or as a native worker on CUDA. Open-Jev-27B-v1.1
+runs as a native Rust/CUDA worker:
| Model | Status |
| --- | --- |
| LAYA | [External worker](recipe/laya/README.md); [Python worker on Apple Silicon (MPS) and CPU](recipe/laya/apple-silicon.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; [H200 validation](recipe/open_jev/validation.md) |
[Supported models and hardware](docs/supported-models.md) lists the devices
and where each worker has been run.
@@ -117,15 +128,16 @@ and where each worker has been run.
## Benchmarks
See the [GPU serving benchmark](benchmarks/README.md) for request replay,
-output-fidelity checks, and the CUDA comparison protocol. GPU performance
-measurements are pending.
+output-fidelity checks, and the CUDA comparison protocol. The
+[Open-Jev H200 results](recipe/open_jev/validation.md) cover 74 single-candidate
+requests and a matched comparison with raw HF Transformers and OpenJev-Fast.
## Roadmap
-The current focus is the Cua-S1 native CUDA worker and the serving benchmark
-harness. Planned work includes the in-repository LAYA model engine, additional
-model engines and GPU backends including Metal, and per-model performance
-measurements as implementations are added and validated.
+The current focus is the native Cua-S1 and Open-Jev CUDA workers and the serving
+benchmark harness. Planned work includes the in-repository LAYA model engine,
+additional model engines and GPU backends including Metal, and per-model
+performance measurements as implementations are added and validated.
--`, then
+ `numactl --membind=0 --physcpubind=0-15`. Run raw HF, native and Fast,
+ preserving the 74-case request order, with one excluded feasibility pass
+ and exactly two measured passes per configuration.
+
+Raw commands, request/token IDs, timing rows, source snapshots and 48 verified
+input hashes are archived locally in the benchmark worktree's
+`profile/jev-hf-transformers-comparison-20261003/`, outside this PR.
+A collector cleanup assertion rejected native's intentional SIGTERM exit after
+all HF/native measurements were saved. Only the remaining Fast configuration
+continued on a second reservation of the same GPU and affinity; no measured
+passes were repeated or added. All task-owned processes exited and GPU 2 returned
+to 0 MB used. The October 2 timings below are a separate experiment.
+
+## Native CUDA Graph replay, 2026-10-03
+
+Warm graph replay reduces mean HTTP latency from **48.086 to 47.112 ms (2.03%)**
+on the same 74-case workload. The shared backend now retains up to 64 exact-length
+graphs, enough for these 57 distinct lengths. Its previous eight-entry cache
+regressed the mixed workload because evicted lengths require another eager
+forward and graph capture. Graph mode remains opt-in with `CUA_S1_GRAPH=1`.
+
+| Workload | Configuration | Mean (ms) | Mean, pass 1 / pass 2 (ms) |
+| --- | --- | ---: | ---: |
+| Fixed 107-token request | Eager | 19.829 | 19.835 / 19.823 |
+| Fixed 107-token request | Graph, eight entries | 18.882 | 18.902 / 18.863 |
+| Fixed 107-token request | Graph, 64 entries | 19.081 | 18.968 / 19.194 |
+| 74 mixed-length requests | Eager | 48.086 | 48.062 / 48.110 |
+| 74 mixed-length requests | Graph, eight entries | 95.289 | 95.425 / 95.153 |
+| 74 mixed-length requests | Graph, 64 entries | 47.112 | 47.050 / 47.174 |
+
+All 74 probabilities and decisions remain exactly unchanged (maximum delta 0.0).
+The 64-entry candidate also reduces the fixed-short mean by 3.77%. Both measured
+passes improve over eager for both workloads, satisfying the prespecified 3%
+short and 2% mixed mean gates. The mixed improvement only narrowly exceeds its
+gate; two passes are observed variability, not confidence intervals or evidence
+for a general workload winner. Observed device memory after mixed passes is
+50,947 / 50,967 / 51,089 MB for eager / graph-eight / graph-64: 122 MB more for
+64 entries than eight. These are scheduler samples, not peak-memory measurements.
+
+This is a separate native A/B experiment, not a newly measured HF/Fast comparison.
+It keeps H200 GPU 2, UUID, NUMA affinity, BF16, max length 16384, model export,
+CUDA library, frontend and request order fixed. Both graph capacities are built
+with identical rustc options and the frozen `202c0e1` dependency artifacts;
+their source copies differ only in the cache limit. Eager uses the eight-entry
+binary with graph mode disabled. Each configuration validates its first long
+inference after readiness to allocate the workload's maximum scratch size, then
+validates the short request. One server is reused for an excluded 32-request
+short feasibility pass and two measured 32-request passes, followed by an excluded
+74-case feasibility pass and two measured 74-case passes. HTTP timing has no
+Nsight launcher and uses the same boundary as the raw HF comparison above.
+
+New-length capture cost remains significant: excluded mixed feasibility means
+are 50.130 / 97.694 / 110.255 ms for eager / graph-eight / graph-64. These are
+single feasibility observations, not measured cold-latency comparisons. Scratch
+growth invalidates graphs, and more than 64 distinct lengths can still evict them.
+No model downloads, cache drops or clock changes occur. All owned processes exit
+and GPU 2 returns to 0 MB used.
+
+A preceding experiment proves warm replay in Nsight Systems: one graph launch,
+zero recaptures and zero individual runtime kernel-launch calls per short request,
+with the same 834 kernels. Node-level graph tracing reports larger gaps despite
+lower unprofiled HTTP latency; use unprofiled measurements for the speedup.
+[NVIDIA documents graph-node tracing overhead](https://docs.nvidia.com/nsight-systems/UserGuide/index.html#cuda-graph-trace).
+Hardware counters and per-SM utilization remain unmeasured.
+
+Raw plans, source copies, build commands, requests/responses, traces, memory
+samples and verified input hashes are archived outside this PR in
+`profile/jev-cuda-graph-20261003-074613/` and
+`profile/jev-cuda-graph-cache64-20261003-075238/`. Run the native worker with
+`CUA_S1_GRAPH=0/1` under the same reservation and affinity documented above;
+warm the maximum workload length and all tested lengths before measured passes.
+For new comparisons, declare the budget first and retain capture costs whenever
+they occur inside measured requests.
+
+## Packed-SiLU A/B, 2026-10-02
+
+A matched comparison of 74 real JevBench `noul` requests, each with one candidate,
+measured packed MLP SiLU with cached residual RMSNorm fixed in both native variants.
+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. Final integration also passed the six-test shared CUDA ABI 4 suite.
+
+### Frozen controls and reproduction
+
+- Device: exact scheduler GPU 2, UUID `GPU-cbf66259-f4ab-0ede-1811-82037dde5924`,
+ NVIDIA H200, 143771 MiB (reported as `NVIDIA L20X` in the archived device metadata);
+ 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 --`,
+then bind the command with `numactl --membind=0 --physcpubind=0-15`. Run baseline,
+packed native and Fast once each, reusing each server for one excluded feasibility
+and two measured passes. Separately trace two requests at each 107/936/3399 tokens
+per configuration with Nsight Systems CUDA/node tracing: 18 traces total.
+
+Raw bodies, timing rows, SHA256 manifests, reproduction commands,
+18 `.nsys-rep`/SQLite pairs and plots are archived locally in the benchmark
+worktree's `profile/jev-single-candidate-silu-pack8-20261002/`, outside this PR.
+Nsight Compute counters are denied by the host policy. Geometry,
+register/shared-memory metadata and CUDA timelines are available; achieved
+occupancy, per-SM tails, stalls, Tensor Core utilization and bandwidth/cache
+efficiency are unmeasured. No hardware-cause claim follows from timeline data alone.
diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md
index b84f50b..133f90e 100644
--- a/src/backends/cuda/qwen3_5/README.md
+++ b/src/backends/cuda/qwen3_5/README.md
@@ -1,9 +1,15 @@
-# Qwen3.5 prefill operations
+# Qwen3.5/3.8 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 current reference tests, including fused gating, cached residual RMSNorm and packed SiLU, passed on H200 (sm_90).
+
+`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 4, which
+includes the CUDA Graph entry points and gated attention.
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/elementwise.cu b/src/backends/cuda/qwen3_5/elementwise.cu
index 3b1c01b..97ae9c5 100644
--- a/src/backends/cuda/qwen3_5/elementwise.cu
+++ b/src/backends/cuda/qwen3_5/elementwise.cu
@@ -72,6 +72,20 @@ __global__ void silu_mul_kernel(const bf16* __restrict__ gate_up, int ld, bf16*
out[i] = to_bf16(round_bf16(silu(f32(row[j]))) * f32(row[I + j]));
}
+__global__ void silu_mul_packed_kernel(const Pack8* __restrict__ gate_up, int ld,
+ Pack8* __restrict__ out, int I, size_t n) {
+ const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
+ if (i >= n) return;
+ const size_t t = i / I, j = i % I;
+ const Pack8* row = gate_up + t * ld;
+ const Pack8 gate = row[j], up = row[I + j];
+ Pack8 result;
+#pragma unroll
+ for (int e = 0; e < 8; e++)
+ result.v[e] = to_bf16(round_bf16(silu(f32(gate.v[e]))) * f32(up.v[e]));
+ out[i] = result;
+}
+
unsigned blocks(size_t n) { return (unsigned)((n + THREADS - 1) / THREADS); }
} // namespace
@@ -120,7 +134,13 @@ extern "C" int cs1_silu_mul(const void* gate_up, int ld, void* out, int T, int I
if (T < 0 || I < 0 || ld < 2 * I) return cudaErrorInvalidValue;
const size_t n = (size_t)T * I;
if (n == 0) return cudaSuccess;
- silu_mul_kernel<<(stream)>>>(
- static_cast(gate_up), ld, static_cast(out), I, n);
+ if (I % 8 == 0 && ld % 8 == 0 &&
+ ((reinterpret_cast(gate_up) | reinterpret_cast(out)) & 15) == 0) {
+ silu_mul_packed_kernel<<(stream)>>>(
+ static_cast(gate_up), ld / 8, static_cast(out), I / 8, n / 8);
+ } else {
+ silu_mul_kernel<<(stream)>>>(
+ static_cast(gate_up), ld, static_cast(out), I, n);
+ }
return cudaGetLastError();
}
diff --git a/src/backends/cuda/qwen3_5/norm.cu b/src/backends/cuda/qwen3_5/norm.cu
index 38d6f98..5e1c210 100644
--- a/src/backends/cuda/qwen3_5/norm.cu
+++ b/src/backends/cuda/qwen3_5/norm.cu
@@ -46,6 +46,32 @@ __global__ void __launch_bounds__(NORM_THREADS)
out[i] = to_bf16(f32(residual[i]) * inv * (1.f + f32(w[i])));
}
+// Keep each thread's rounded residuals in registers across the reduction. Retain
+// the scalar kernel's element assignment, sum order, and 256-thread block.
+template
+__global__ void __launch_bounds__(NORM_THREADS)
+ add_rms_norm_cached_kernel(bf16* __restrict__ residual, const bf16* __restrict__ delta,
+ const bf16* __restrict__ w, bf16* __restrict__ out, float eps) {
+ __shared__ float scratch[32];
+ const size_t row = (size_t)blockIdx.x * D;
+ float r[D / NORM_THREADS];
+ float ss = 0.f;
+#pragma unroll
+ for (int i = 0; i < D / NORM_THREADS; i++) {
+ const int col = threadIdx.x + i * NORM_THREADS;
+ const bf16 value = to_bf16(f32(residual[row + col]) + f32(delta[row + col]));
+ residual[row + col] = value;
+ r[i] = f32(value);
+ ss += r[i] * r[i];
+ }
+ const float inv = rsqrtf(block_sum(ss, scratch) / D + eps);
+#pragma unroll
+ for (int i = 0; i < D / NORM_THREADS; i++) {
+ const int col = threadIdx.x + i * NORM_THREADS;
+ out[row + col] = to_bf16(r[i] * inv * (1.f + f32(w[col])));
+ }
+}
+
// Qwen3_5RMSNormGated with D = 128, one warp per row: the normalized value is
// rounded to bfloat16, multiplied by w in bfloat16, then by silu(z) in float32.
__global__ void gated_rms_norm_kernel(const bf16* __restrict__ x, const bf16* __restrict__ z, int ldz,
@@ -87,6 +113,17 @@ extern "C" int cs1_rms_norm(const void* x, const void* w, void* out, int rows, i
extern "C" int cs1_add_rms_norm(void* residual, const void* delta, const void* w, void* out, int rows, int D,
float eps, void* stream) {
if (rows <= 0) return cudaSuccess;
+ switch (D) {
+#define CACHED_NORM(width) \
+ case width: \
+ add_rms_norm_cached_kernel<<(stream)>>>( \
+ static_cast(residual), static_cast(delta), static_cast(w), \
+ static_cast(out), eps); \
+ return cudaGetLastError()
+ CACHED_NORM(2560);
+ CACHED_NORM(5120);
+#undef CACHED_NORM
+ }
add_rms_norm_kernel<<(stream)>>>(
static_cast(residual), static_cast(delta), static_cast(w),
static_cast(out), D, eps);
diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h
index ec0dcd6..ae5ce4d 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 3
+// Bumped whenever the required interface below changes.
+#define CS1_ABI_VERSION 4
#ifdef __cplusplus
extern "C" {
@@ -81,6 +81,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..5384417
--- /dev/null
+++ b/src/models/open_jev/native/Cargo.toml
@@ -0,0 +1,18 @@
+[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"] }
+
+[[test]]
+name = "contract"
+path = "../../../../tests/open_jev/contract.rs"
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..a1ec531
--- /dev/null
+++ b/src/models/open_jev/native/src/contract.rs
@@ -0,0 +1,222 @@
+//! 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,
+}
+
+/// Python json.dumps(ensure_ascii=False, sort_keys=True), or the string itself.
+fn render(value: &Value) -> String {
+ if let Some(text) = value.as_str() {
+ return text.to_owned();
+ }
+ let mut sorted = value.clone();
+ sorted.sort_all_objects();
+ json::dumps(&sorted)
+}
+
+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..29b76f9
--- /dev/null
+++ b/src/models/open_jev/native/src/engine.rs
@@ -0,0 +1,186 @@
+//! 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)]
+#[path = "../../../../../tests/open_jev/tokenization.rs"]
+mod tests;
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/qwen3_5/native/Cargo.toml b/src/models/qwen3_5/native/Cargo.toml
new file mode 100644
index 0000000..b71c29a
--- /dev/null
+++ b/src/models/qwen3_5/native/Cargo.toml
@@ -0,0 +1,23 @@
+[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"] }
+
+[[test]]
+name = "config"
+path = "../../../../tests/qwen3_5/config.rs"
+
+[[test]]
+name = "kernels"
+path = "../../../../tests/qwen3_5/kernels.rs"
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 719dc45..7433ee6 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 = 3;
+const ABI_VERSION: u32 = 4;
pub const LIBRARY: &str = "libqwen3_5_cuda.so";
/// A `cudaStream_t`.
@@ -95,6 +95,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 74%
rename from src/models/cua_s1/native/src/json.rs
rename to src/models/qwen3_5/native/src/json.rs
index a0e2ed5..18529af 100644
--- a/src/models/cua_s1/native/src/json.rs
+++ b/src/models/qwen3_5/native/src/json.rs
@@ -174,72 +174,5 @@ pub fn float_repr(x: f64) -> String {
}
#[cfg(test)]
-mod tests {
- use super::*;
-
- fn err(body: &str) -> String {
- parse(body.as_bytes()).unwrap_err()
- }
-
- #[test]
- fn float_repr_matches_python_examples() {
- let cases = [
- (1.0, "1.0"),
- (1e16, "1e+16"),
- (1e15, "1000000000000000.0"),
- (1e-5, "1e-05"),
- (1e-4, "0.0001"),
- (-0.0, "-0.0"),
- (3.14e-07, "3.14e-07"),
- (5e-324, "5e-324"),
- (1.7976931348623157e308, "1.7976931348623157e+308"),
- (5.960464477539063e-08, "5.960464477539063e-08"),
- ];
- for (x, want) in cases {
- assert_eq!(float_repr(x), want, "{x:e}");
- }
- }
-
- #[test]
- fn dumps_matches_python() {
- let v = Value::Object(parse(br#"{"a": [1.0, 1e16, 1e-5, 0.0001, -0.0, 123456789012345678, 3.14e-07, true, null], "b": {}}"#).unwrap());
- assert_eq!(
- dumps(&v),
- r#"{"a": [1.0, 1e+16, 1e-05, 0.0001, -0.0, 123456789012345678, 3.14e-07, true, null], "b": {}}"#
- );
- let s = Value::String("\u{0}\u{1f}\u{7f}\u{2028}\"\\/\t\u{8}\u{c}é😀".into());
- assert_eq!(
- dumps(&s),
- "\"\\u0000\\u001f\u{7f}\u{2028}\\\"\\\\/\\t\\b\\fé😀\""
- );
- }
-
- #[test]
- fn rejects_what_the_contract_rejects() {
- assert_eq!(err("[]"), "request body must be a JSON object");
- for body in [
- r#"{"a": NaN}"#,
- r#"{"a": 1e400}"#,
- r#"{"a": "\ud800x"}"#,
- r#"{"a": 1, "b": 2, "a": 3}"#,
- r#"{"a": [1,]}"#,
- r#"{} x"#,
- "\u{feff}{}",
- ] {
- assert!(
- err(body).starts_with("request body is not valid JSON"),
- "{body}"
- );
- }
- assert!(err(r#"{"a": 1, "a": 2}"#).contains("duplicate key \"a\""));
- assert!(
- err(&format!(
- "{{\"a\": {}1{}}}",
- "[".repeat(200),
- "]".repeat(200)
- ))
- .contains("recursion limit")
- );
- assert!(parse(b"{\"a\": \"\xff\"}").is_err());
- }
-}
+#[path = "../../../../../tests/qwen3_5/json.rs"]
+mod tests;
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 ae14d55..c8ab5d6 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
@@ -662,7 +662,7 @@ impl Model {
cuda::synchronize(self.stream)?;
match cuda::Graph::capture(self.stream, || self.run(s, t)) {
Ok(graph) => {
- if self.graphs.len() == 8 {
+ if self.graphs.len() == 64 {
self.graphs.pop_front();
}
self.graphs.push_back((t, graph));
@@ -830,11 +830,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,
@@ -843,16 +844,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/tests/open_jev/contract.rs b/tests/open_jev/contract.rs
new file mode 100644
index 0000000..4c63afa
--- /dev/null
+++ b/tests/open_jev/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/tests/open_jev/data/contract.json b/tests/open_jev/data/contract.json
new file mode 100644
index 0000000..e90a417
--- /dev/null
+++ b/tests/open_jev/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/tests/open_jev/data/tokenization.json b/tests/open_jev/data/tokenization.json
new file mode 100644
index 0000000..4dc56ad
--- /dev/null
+++ b/tests/open_jev/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/tests/open_jev/tokenization.rs b/tests/open_jev/tokenization.rs
new file mode 100644
index 0000000..bba7363
--- /dev/null
+++ b/tests/open_jev/tokenization.rs
@@ -0,0 +1,22 @@
+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!("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/tests/qwen3_5/config.rs b/tests/qwen3_5/config.rs
new file mode 100644
index 0000000..b7706ef
--- /dev/null
+++ b/tests/qwen3_5/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/qwen3_5/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/tests/qwen3_5/data/config.json b/tests/qwen3_5/data/config.json
new file mode 100644
index 0000000..706cebd
--- /dev/null
+++ b/tests/qwen3_5/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/tests/qwen3_5/json.rs b/tests/qwen3_5/json.rs
new file mode 100644
index 0000000..4a1edc5
--- /dev/null
+++ b/tests/qwen3_5/json.rs
@@ -0,0 +1,67 @@
+use super::*;
+
+fn err(body: &str) -> String {
+ parse(body.as_bytes()).unwrap_err()
+}
+
+#[test]
+fn float_repr_matches_python_examples() {
+ let cases = [
+ (1.0, "1.0"),
+ (1e16, "1e+16"),
+ (1e15, "1000000000000000.0"),
+ (1e-5, "1e-05"),
+ (1e-4, "0.0001"),
+ (-0.0, "-0.0"),
+ (3.14e-07, "3.14e-07"),
+ (5e-324, "5e-324"),
+ (1.7976931348623157e308, "1.7976931348623157e+308"),
+ (5.960464477539063e-08, "5.960464477539063e-08"),
+ ];
+ for (x, want) in cases {
+ assert_eq!(float_repr(x), want, "{x:e}");
+ }
+}
+
+#[test]
+fn dumps_matches_python() {
+ let v = Value::Object(parse(br#"{"a": [1.0, 1e16, 1e-5, 0.0001, -0.0, 123456789012345678, 3.14e-07, true, null], "b": {}}"#).unwrap());
+ assert_eq!(
+ dumps(&v),
+ r#"{"a": [1.0, 1e+16, 1e-05, 0.0001, -0.0, 123456789012345678, 3.14e-07, true, null], "b": {}}"#
+ );
+ let s = Value::String("\u{0}\u{1f}\u{7f}\u{2028}\"\\/\t\u{8}\u{c}é😀".into());
+ assert_eq!(
+ dumps(&s),
+ "\"\\u0000\\u001f\u{7f}\u{2028}\\\"\\\\/\\t\\b\\fé😀\""
+ );
+}
+
+#[test]
+fn rejects_what_the_contract_rejects() {
+ assert_eq!(err("[]"), "request body must be a JSON object");
+ for body in [
+ r#"{"a": NaN}"#,
+ r#"{"a": 1e400}"#,
+ r#"{"a": "\ud800x"}"#,
+ r#"{"a": 1, "b": 2, "a": 3}"#,
+ r#"{"a": [1,]}"#,
+ r#"{} x"#,
+ "\u{feff}{}",
+ ] {
+ assert!(
+ err(body).starts_with("request body is not valid JSON"),
+ "{body}"
+ );
+ }
+ assert!(err(r#"{"a": 1, "a": 2}"#).contains("duplicate key \"a\""));
+ assert!(
+ err(&format!(
+ "{{\"a\": {}1{}}}",
+ "[".repeat(200),
+ "]".repeat(200)
+ ))
+ .contains("recursion limit")
+ );
+ assert!(parse(b"{\"a\": \"\xff\"}").is_err());
+}
diff --git a/src/models/cua_s1/native/tests/kernels.rs b/tests/qwen3_5/kernels.rs
similarity index 54%
rename from src/models/cua_s1/native/tests/kernels.rs
rename to tests/qwen3_5/kernels.rs
index 692f1b5..469244e 100644
--- a/src/models/cua_s1/native/tests/kernels.rs
+++ b/tests/qwen3_5/kernels.rs
@@ -3,12 +3,12 @@
//! asked for:
//!
//! 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
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,142 @@ fn from_device(buf: &DeviceBuffer, n: usize, st: Stream) -> Vec {
.collect()
}
+#[test]
+#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"]
+fn residual_rms_norm_matches_rounded_reference() {
+ let st = setup();
+ let eps = 1e-6f32;
+ // Cua-S1 and Open-Jev widths, scalar fallbacks for other widths, and a BF16
+ // element offset to check that cached loads need no packed alignment.
+ for (d, offset) in [
+ (2560usize, 0usize),
+ (5120, 0),
+ (8192, 0),
+ (257, 0),
+ (8448, 0),
+ (5120, 1),
+ ] {
+ for rows in [1usize, 3, 107] {
+ let n = rows * d;
+ let mut residual = random(n + offset, 5, 3.0);
+ let mut delta = random(n + offset, 6, 0.5);
+ let mut weights = random(d + offset, 7, 0.25);
+ // Exercise the BF16 rounding between addition and normalization.
+ residual[offset] = bf16::ONE;
+ delta[offset] = bf16::from_f32(1.0 / 256.0);
+ weights[offset + 1] = bf16::from_f32(-1.0);
+ let rounded: Vec = residual[offset..]
+ .iter()
+ .zip(&delta[offset..])
+ .map(|(r, d)| bf16::from_f32(r.to_f32() + d.to_f32()))
+ .collect();
+ let input = to_device(&residual, st);
+ let change = to_device(&delta, st);
+ let weight = to_device(&weights, st);
+ let output = DeviceBuffer::new((n + offset) * 2).unwrap();
+ // SAFETY: complete BF16 rows and d weights after the optional one-
+ // element offset, with separate input, delta, weight, and output.
+ unsafe {
+ check(
+ (api().cs1_add_rms_norm)(
+ input.at(offset * 2),
+ change.at(offset * 2),
+ weight.at(offset * 2),
+ output.at(offset * 2),
+ rows as i32,
+ d as i32,
+ eps,
+ st,
+ ),
+ "residual norm",
+ )
+ .unwrap();
+ }
+ let added = from_device(&input, n + offset, st);
+ assert_eq!(
+ added[offset..],
+ rounded.iter().map(|v| v.to_f32()).collect::>(),
+ "residual addition changed BF16 rounding at d={d}, offset={offset}"
+ );
+ if offset != 0 {
+ assert_eq!(added[0], residual[0].to_f32());
+ }
+ let got = from_device(&output, n + offset, st);
+ for (row, values) in rounded.chunks_exact(d).enumerate() {
+ let ss: f64 = values.iter().map(|v| (v.to_f32() as f64).powi(2)).sum();
+ let inv = ((ss / d as f64) as f32 + eps).sqrt().recip();
+ for (i, value) in values.iter().enumerate() {
+ let want =
+ bf16::from_f32(value.to_f32() * inv * (1.0 + weights[offset + i].to_f32()));
+ let actual = bf16::from_f32(got[offset + row * d + i]);
+ assert!(actual.is_finite());
+ // Reduction order and CUDA rsqrt may move a value by one
+ // BF16 ULP; residual addition above must remain exact.
+ assert!(
+ actual.to_bits().abs_diff(want.to_bits()) <= 1,
+ "norm mismatch d={d}, offset={offset}, row={row}, i={i}: {actual} vs {want}"
+ );
+ }
+ }
+ }
+ }
+}
+
+#[test]
+#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"]
+fn silu_mul_matches_rounded_reference() {
+ let st = setup();
+ // Packed model widths, padded rows, and scalar fallbacks for an odd width
+ // or stride. A one-element pointer offset forces the scalar path on the
+ // same inputs, which must agree exactly with packed execution.
+ for (width, ld) in [(8192usize, 16384usize), (17408, 34816), (64, 136), (37, 79)] {
+ for rows in [1usize, 3, 107] {
+ let input = random(rows * ld, 23, 12.0);
+ let mut results = Vec::new();
+ for offset in [0usize, 1] {
+ let mut padded = vec![bf16::ZERO; offset];
+ padded.extend_from_slice(&input);
+ let source = to_device(&padded, st);
+ let output = DeviceBuffer::new((rows * width + offset) * 2).unwrap();
+ // SAFETY: complete gate/up rows and an output of rows*width
+ // BF16 elements after the optional one-element offset.
+ unsafe {
+ check(
+ (api().cs1_silu_mul)(
+ source.at(offset * 2),
+ ld as i32,
+ output.at(offset * 2),
+ rows as i32,
+ width as i32,
+ st,
+ ),
+ "silu_mul",
+ )
+ .unwrap();
+ }
+ results.push(from_device(&output, rows * width + offset, st)[offset..].to_vec());
+ }
+ assert_eq!(
+ results[0], results[1],
+ "packed/scalar mismatch at {width}/{ld}"
+ );
+ for (row, values) in input.chunks_exact(ld).enumerate() {
+ for j in 0..width {
+ let gate = values[j].to_f32() as f64;
+ let activated = bf16::from_f64(gate / (1.0 + (-gate).exp()));
+ let want = bf16::from_f32(activated.to_f32() * values[width + j].to_f32());
+ let actual = bf16::from_f32(results[0][row * width + j]);
+ assert!(actual.is_finite());
+ assert!(
+ actual.to_bits().abs_diff(want.to_bits()) <= 1,
+ "SiLU rounding mismatch at {width}/{ld}, row={row}, j={j}: {actual} vs {want}"
+ );
+ }
+ }
+ }
+ }
+}
+
#[test]
#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"]
fn graph_replay_reads_updated_inputs_after_failed_capture() {
@@ -131,6 +267,126 @@ fn graph_replay_reads_updated_inputs_after_failed_capture() {
}
}
+#[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() {
@@ -270,7 +526,23 @@ fn gated_delta_reference(
fn gated_delta_rule_matches_recurrent_reference() {
let st = setup();
let (h, hk, d) = (4usize, 2usize, 128usize);
- for t in [1usize, 64, 150] {
+ // Exercise inverse tile/chunk boundaries, the short model request, and
+ // accumulated state with weak decay over many chunks.
+ for (t, decay_scale) in [
+ (1usize, 1.0f32),
+ (16, 1.0),
+ (32, 1.0),
+ (63, 1.0),
+ (64, 1.0),
+ (65, 1.0),
+ (107, 1.0),
+ (127, 1.0),
+ (128, 1.0),
+ (129, 1.0),
+ (150, 1.0),
+ (107, 0.01),
+ (936, 0.01),
+ ] {
// q close to k, so that q.k and the outputs are of order one as in the model
let k = random(t * hk * d, 12, 1.0);
let q: Vec = k
@@ -282,7 +554,7 @@ fn gated_delta_rule_matches_recurrent_reference() {
// log decays in (-2, 0) and learning rates in (0, 1), as sigmoid and -exp * softplus give
let g: Vec = random(t * h, 14, 1.0)
.iter()
- .map(|x| x.to_f32() - 1.0)
+ .map(|x| (x.to_f32() - 1.0) * decay_scale)
.collect();
let beta: Vec = random(t * h, 15, 0.5)
.iter()