[RFC]: Native Rust/CUDA support for MaestroYan/JEMM and community help #121

Motivation

Add support for MaestroYan/JEMM, an
open-weight Apache-2.0 decision model advertised as "Like Jev, but multimodal
and open-weight." One prefill-only forward pass turns a state (text plus up to
four screenshots), a question, and 2–32 candidates into a probability for every
candidate, for choice, noul, and score questions.

JEMM is a strong fit for this engine and unusually cheap to integrate relative
to the rest of the Jev family:

• It is a rank-16 LoRA adapter on Qwen/Qwen3.8-27B — the same backbone our
  validated Open-Jev-27B native worker already runs on (shared
  Qwen3.5/3.8 prefill kernels, merged-checkpoint export machinery).
• Its readout is Cua-S1-style last-position label-token logits over 32
  labels (A–Z, 0–5) with a saved temperature — no trained decision
  head, no full-vocabulary features, no branching (unlike Open-Jev-27B,
  openjev/openjev [RFC]: Support openjev/openjev in System1-Omni — community help wanted #95, or OmniJev [RFC]: Native Rust/CUDA support for OmniJev-4B v1.1 and community help #114).
• Its reference server (ypcypc/JEMM) already
  speaks the Jev-style POST /v1/systemone envelope our frontend forwards, and
  is serial by construction, matching our shared serial runtime.
• It is multimodal (Qwen3_5ForConditionalGeneration), and main now has
  shared CUDA vision entry points (ABI 5) plus native Cua-S1 vision and image
  preprocessing in Rust to reuse.

No existing worker can load this checkpoint: Open-Jev-27B's worker compiles one
independent prompt per candidate and scores it with a trained scalar head, while
JEMM renders all candidates into one prompt and reads 32 label tokens at the
final position. This is a separate model directory and recipe, not a config
change.

This RFC invites feedback and contributors; nothing below has been implemented.

Proposed Change

1. Target and reproducible reference

The adapter's adapter_config.json names Qwen/Qwen3.8-27B but does not pin a
base revision ("revision": null); the pin above is proposed because it is the
revision our Open-Jev-27B validation already used. Confirm with the authors and
freeze a known-good environment before generating fixtures. The server's
declared environment is torch>=2.12, transformers>=5.17,<6, peft>=0.21,
accelerate>=1.15, flash-linear-attention>=0.5.2 (Linux), numpy, pillow
— minimum versions, not a lock.

The adapter is LoRA rank 16, alpha 32, dropout 0, stored fp32 by PEFT. Its
targets cover attention and MLP projections and the Gated DeltaNet
projections (in_proj_a, in_proj_b, in_proj_qkv, in_proj_z,
out_proj) — unlike the Cua-S1 text adapter, which leaves the linear-attention
blocks alone. The export must merge all of these.

decision_config.json holds three scalars the worker must load, not hard-code:
temperature = 1.3480874159655591 (text), mm_temperature = 1.3954832341582943
(any image present), threshold = 0.9872681877423998.

First slice: BF16 CUDA, text and screenshots (1–4 images per request), all
three question types, 2–32 candidates, serial admission, served behind the
existing frontend. Out of scope initially: the reference's multi-question
"bundle" prompt, cross-question prefix reuse, quantization, Metal, multi-GPU.

2. Ownership and serving interface

Follow the architecture contracts and the model integration skill:

• src/models/jemm/ owns the native worker: request compilation, candidate
  rendering, tokenization, image preparation, execution, readout, and response
  finishing; recipe/jemm/ holds export/setup/launch and example requests; all
  tests and fixtures live in tests/jemm/ under the root tests/ tree.
• Serve POST /v1/systemone and GET /health behind the existing Rust
  frontend unchanged (OMNI_JEV_BACKEND_URL pointing at the JEMM worker).
• Wire input matches the reference server: state (string, or any JSON value
  rendered as compact JSON), questions as an object keyed by question ID, each
  with type (choice/noul/score, defaulting to choice) and criteria,
  plus an optional images list of up to 4 base64 PNG/JPEG/WebP data URLs. Any
  transport adaptation must be documented and tested.
• Reuse the shared serial runtime for admission. The reference server is serial
  behind a global lock; a native worker must preserve single-request decision
  semantics and document its concurrency model.
• The serving process requires no Python/PyTorch. A one-time exporter and a
  pinned Python reference for fixtures are appropriate.

3. Checkpoint preparation and execution reuse

• Reuse the Open-Jev export pattern (recipe/open_jev/export_merged.py): merge
  the LoRA into the 27B backbone in BF16 (CPU, roughly the same ~110 GB RAM /
  ~52 GB storage class as Open-Jev-27B, plus vision weights), and save the
  decision_config scalars, tokenizer/processor files, and chat-template
  arguments in a manifest (e.g. jemm_export.json). The worker refuses a plain
  base checkpoint or an incomplete export.
• Qualify BF16 LoRA merging against the unmerged fp32-PEFT reference before
  adopting it (same concern as [RFC]: Native Rust/CUDA support for OmniJev-4B v1.1 and community help #114). The merge must cover the GDN projections
  listed above.
• Verify input/output embedding tying for this checkpoint from its config
  before reading label rows (see the [RFC]: Support openjev/openjev in System1-Omni — community help wanted #95 warning about untied heads in this
  family).
• Reuse the shared Qwen3.5/3.8 prefill kernels and layer loop
  (src/models/qwen3_5/native/, src/backends/cuda/qwen3_5/). For images,
  reuse the shared vision entry points and the native Cua-S1 vision /
  image-preprocessing pipeline where the 27B vision tower configuration
  permits; CPU image preparation must match the base AutoProcessor exactly
  (decode, RGB, resize, patch order, placeholder expansion, 3D positions).
  Existing Cua-S1, Open-Jev-27B, and Open-Jev-9B regression checks must remain
  intact; any ABI change bumps the version and rebuilds consumers together.
• GPU requirement follows the backbone: BF16 27B plus vision, one CUDA GPU with
  ≥64 GB as the reference states; compute capability 8.0+ per the shared
  kernels.

4. JEMM semantics to preserve

The pinned normative reference is the server code (ypcypc/JEMM@6822fe0),
not the model-card snippet: where they differ, the server wins, and known
differences must be recorded.

5. Execution and optimization plan

The reference re-encodes the full state+images prompt per question. The
obvious native win is encoding the multimodal prefix once per request and
running question suffixes against it — exactly the prefix-continuation work in
flight (#97, #99, and jev-vl #96's prepared image embeddings + language-prefix
state). The reference's own bundle mode is a third formulation of the same
idea.

Proposed sequence:

1. Contract/export: pinned environment, artifact inventory, CPU fixtures
   for every boundary in §4, documented wire input/output.
2. Native eager worker: full text+vision execution, one forward per
   question, serial admission, frontend integration, parity vs the pinned
   reference server.
3. Request-local prefix reuse: shared multimodal prefix across questions of
   one request, coordinated with jev_vl: add native inference and prefix caching #96/qwen3_5: add prefix continuation operations #97/qwen3_5: run prompts that share prefixes once per prefix #99/decider: reuse request prefixes with conservative auto selection #113; the bundle readout is an
   alternative formulation to evaluate with fixtures against the reference's
   encode_bundle.
4. Further optimization: CUDA Graphs / packed GEMMs where profiling
   supports them, coordinated with perf(open-jev): pack candidate projections within requests #102/cua_s1: add native multimodal language CUDA Graph replay #103/cua_s1: add native vision CUDA Graph replay #108/feat(decider): independent CUDA Graph replay mode #112.

Each stage is a separate hypothesis measured against its own baseline; no
combined speedup claims.

6. Validation and acceptance

• Freeze code/artifact hashes, dependency versions, dtypes, and numerical
  gates before GPU work. Suggested initial gates (proposals, not results):
  max absolute per-candidate probability drift 0.01 vs the pinned reference
  server; matching decisions when the reference top-two margin exceeds
  0.02; report near-threshold confidence crossings explicitly.
• CPU fixtures (exact, from the pinned reference): rendered prompts and
  token IDs for all three question types, default type, criteria ordering,
  tool-spec rendering, whitespace-collapse vs strip() distinction, compact
  JSON state, Unicode and special-token text, 32 label token IDs, image
  decode/pixel-cap/format errors, budget and body-size rejections, response
  finishing including ties and expected_value, usage accounting, and
  malformed inputs. All tests/helpers/fixtures under root tests/jemm/.
• Reserved-GPU parity vs the pinned reference server on the same
  checkpoint: text and image requests, every question type, staged comparisons
  (merged-export vs unmerged-PEFT; native vision outputs vs processor
  outputs; final probabilities and decisions).
• Direct-worker and frontend-proxied HTTP checks, including the first
  request after readiness and health behavior. Revalidate Cua-S1, Open-Jev-27B,
  and Open-Jev-9B consumers when shared code or the CUDA ABI changes.
• Performance runs under the benchmark protocol
  with a JEMM-aware response validator (32-candidate probability sums; bench: accept wire-rounded probability sums per label count #75's
  per-label-count rounding tolerance already covers this shape). Separate
  preparation, startup, first inference, warm latency, and throughput; retain
  raw results.
• Update recipe/jemm/, src/models/jemm/README.md,
  docs/supported-models.md, root README, mkdocs.yml, and tracker [Tracking]: Model support status #83 when
  validated support lands. Record tested GPUs and scope; no general claims
  from compilation or narrow fixtures.

Alternatives considered

• External Python worker first: the reference server already exposes
  /v1/systemone and /health, so it can run behind the existing frontend
  (OMNI_JEV_BACKEND_URL) with zero code, like the LAYA external-worker
  recipe. Useful as an immediate recipe and oracle; it does not meet the
  native-serving goal and its throughput is bounded by the per-question
  re-encoding.
• Extending the Open-Jev worker: rejected; different prompt compilation,
  readout, calibration, and modality would fork that model's validated
  contract.
• Adopting the bundle readout first: higher ceiling per request, but it
  front-loads the hardest readout semantics before basic parity exists.

Community help wanted

Comment with the task you want, relevant experience, and available hardware.
Small PRs and design review are welcome; nobody needs to own the whole
integration.

Review questions:

1. Is text+vision in the first native slice right, or text-first with vision
   second (given the shared vision entry points and Cua-S1 precedent)?
2. Server (float64) or card (torch float32) softmax as the normative
   calibration reference?
3. Should the worker surface an explicit undecided signal from threshold, or
   preserve the reference server's probabilities+confidence shape exactly?
4. Which prefix-reuse formulation should JEMM adopt — question suffixes on a
   shared multimodal prefix (jev_vl: add native inference and prefix caching #96/qwen3_5: add prefix continuation operations #97/qwen3_5: run prompts that share prefixes once per prefix #99) or the reference's bundle readout?

Feedback Period

Initial feedback through 2026-10-22; contributions and discussion remain
welcome afterward.

CC List

@MaestroYan — the model author; confirmation of the base pin, environment, and
training-time image assumptions would be especially helpful.
@ypcypc — the reference server author; contract clarifications (undecided
semantics, bundle plans) likewise.

Anything else

Preparation for this RFC inspected public model metadata and the pinned
reference source only. No model weights were loaded, no GPU inference or
benchmarks were run, and no native speedup is claimed.

JEMM is one of several Jev-family integrations now in flight; it is a distinct
checkpoint from the supported Open-Jev-27B/9B (ZefanCai), openjev/openjev
(#95), autotrust/JEV-27B-VL (#96), and OmniJev (#114), and is not affiliated
with TypeSafe AI.

License notes: adapter and reference server are Apache-2.0; the card's training
mix includes CC-BY-3.0/CC-BY-4.0 and OpenRAIL components (attribution only —
no weights are distributed here); confirm the base checkpoint's license at pin
time.

Related: #9 (roadmap), #61 (shared /v1/systemone contract), #83 (model
tracker), #85 (observation reuse), #95/#114 (sibling Jev-family RFCs), #96
(je-vl native + prefix caching), #97/#99 (prefix continuation/sharing), #102
(packed prefill, merged), #103/#108 (native multimodal graphs).

Before submitting a new issue

• Searched existing issues and PRs for MaestroYan/JEMM, JEMM, and the
  checkpoint ID; reviewed the RFC template, architecture contracts, supported
  -model documentation, current native code, and the pinned upstream reference.
  No existing JEMM integration was found.