[Model] Add Laya typed-decision models (Jev /v1/systemone via structured_server)

[LucasWilkinson]

Overview

Adds Laya typed-decision models (convaiinnovations/laya, convaiinnovations/laya-multilingual) as a token_classify pooling model, LayaForDecision. It also lets the #57250 structured-decisions server (/v1/systemone, Jev's API) front Laya with --backend laya.

Laya came up on #57250 as an encoder alternative to DiffusionGemma structured reads. It answers choice / score / noul questions about a state in one forward pass: a ModernBERT/mmBERT encoder plus a small decision head, with no generation.

Claims

• Model support. vLLM serves both public Laya checkpoints (after a one-time layout conversion) and reproduces Laya's own Agent. Prompt token ids are identical, every answer matches, and probabilities agree to rounding in fp32.
• Jev API. structured_server.py --backend laya serves Jev's /v1/systemone and the schema-in-system-message /v1/chat/completions in front of vLLM's /pooling. It supports depends_on, ask_if, ask and instructions. think, images and raw chat are rejected, because Laya does not generate text.
• No behavior change for DiffusionGemma. The two helpers factored out of the diffusion path (answer_from, skip_reason) leave its output byte-identical.
• Throughput. Throughput is on par with Laya's reference runtime. This is not a speedup claim, see Validation.

Validation

All on B300, rebased on cb4f016c5d (model runner v2 active).

Parity against Laya's reference Agent, using a script in the Laya repo's layout: 4 states (dict, text, conversation list, Hindi) × 6 questions (4-way choice, score, noul, noul with custom labels, binary choice, 12-way choice) per checkpoint.

Check	Result
Prompt token ids vs laya.common.build_sequence	identical, 24/24 per checkpoint
Chosen answer	same in 48/48
fp32 + FLEX_ATTENTION, TF32 off	worst probability gap 1e-4 (the 4-decimal rounding of the answers)
bf16 default (TRITON_ATTN)	worst gap 0.017 (English), 0.028 (multilingual)
Through structured_server.py --backend laya	worst gap 0.0054, same answers

The bf16 gaps are precision noise, not implementation error: Laya runs under autocast with fp32 LayerNorms, while vLLM runs fully in bf16. The fp32/FLEX row isolates this.

CI test tests/models/language/pooling/test_laya.py:

• It builds a small random checkpoint in the converted layout and compares against a PyTorch port of Laya's DecisionModel. One batch mixes question types, option counts (2 to 12, which hits a temperature bucket) and a prompt without markers.
• The public checkpoints need conversion, so CI cannot load them directly.
• Mutation check: a wrong question type, an ignored temperature bucket, a wrong [CLS] offset and off-by-one markers each fail the test by 0.036 to 0.62. The real error is 1.5e-3 and the tolerance is 1e-2, so the TF32 attention gap leaves margin on other GPUs.
• The file sits in tests/models/language/pooling/, which the Language Models (Extra Pooling) step runs.

pytest tests/models/language/pooling/test_laya.py   # 1 passed

End to end:

python examples/pooling/token_classify/laya/convert_checkpoint.py convaiinnovations/laya ./laya-vllm
vllm serve ./laya-vllm --served-model-name laya
python examples/features/structured_diffusion/structured_server.py --backend laya \
    --upstream http://127.0.0.1:8000 --model laya --tokenizer ./laya-vllm

• Basic /v1/systemone, ask_if + depends_on staging, the chat-completions schema route and the rejection paths behave as described.
• A prompt with no markers or a bad question type returns an empty output, and the engine stays healthy.
• laya_offline.py runs.

DiffusionGemma path: the pre- and post-PR structured_server.decide() were run with the upstream calls stubbed, across samples = 3 / auto / 1 and a schema using ask_if, depends_on and alone. The outputs are byte-identical. There is no DiffusionGemma checkpoint on this machine, so the diffusion path was not run end to end.

Throughput, English checkpoint, 6 questions per request:

	decisions/s
vLLM, through the HTTP proxy (Python tokenization), concurrency 32 / 128	1522 / 1624
Laya predict_batch, batch 8 / 32	1291 / 1554

A single 6-question request takes 14 ms p50. These numbers are from the pre-rebase build.

Details

Design.

• LayaForDecision reuses ModernBertModel unchanged. On top it adds Laya's question-type embedding and its two pre-norm head layers, the equivalent of nn.TransformerEncoderLayer(norm_first=True), running on EncoderOnlyAttention.
• The question type is read from the token after each sequence's [CLS] (where positions == 0).
• LayaDecisionPooler finds each option's [MASK] marker from the prompt token ids and scores it. It applies Laya's per-type and per-option-count temperatures, clamped to [0.5, 5] as Laya's Agent does, and runs the action head on [CLS].
• Output is one row per option, [p_option, *p_action]. use_activation=false returns logits.
• The pooler never raises, because it runs inside the engine core: malformed prompts, including profiling dummies, get an empty output.

Split between model and server, the same split as DiffusionGemma. The model definition owns the readout: marker detection, question type and calibration. Prompt construction and answer decoding stay client-side, because Laya budgets in tokens (options capped at 48 tokens, the header at head_max_len, conversations trimmed from their oldest turn). A chat template cannot express that without changing the token ids. An IO-processor plugin would allow a proxy-free /v1/systemone, but plugins currently load only from installed packages, so that is left for a follow-up.

Checkpoint conversion. The Hugging Face repos keep config.json and the tokenizer in subfolders, so convert_checkpoint.py writes a loadable directory:

• the encoder config with architectures, a laya_config block (head layers, temperatures, mask_token_id, question-type token ids) and dtype set from Laya's amp_dtype;
• the fixed tokenizer;
• the weights, hard-linked.

A repo with a root config.json would remove this step and allow a normal model test. The registry entry is is_available_online=False until then.

Compatibility and limitations.

• Tensor parallel >1 fails, from a pre-existing bug. vLLM's ModernBertAttention splits the TP-sharded QKV output using the global head count (modernbert.py:147), so every ModernBERT model fails at TP>1 on main. This PR does not touch that file. Laya's checkpoints are 322M/421M parameters, so TP is rarely useful. The fix is small (per-rank head counts) and is better as its own PR.
• Chunked prefill and prefix caching are off by default for bidirectional pooling models. Forcing them on is unsupported, with the same warning as other encoder-only models. The question-type lookup relies on this.
• No pipeline parallel. SupportsPP is not implemented.
• Head layers run eagerly, outside torch.compile. The encoder is still compiled and CUDA-graphed.
• Laya's Router (English vs multilingual by script) is not ported. Run one server per checkpoint.
• Saturated action head. In the English checkpoint act_probability is always 1.0. Laya's runtime gives the same.
• Complexity. About 268 new lines under vllm/, almost all in the self-contained laya.py, plus a one-line registry entry. There are no changes to core components.

Duplicate-work check. Open PR and issue searches for laya and related terms (re-run 2026-09-29) found nothing. Laya appears only in a comment on #57250.

AI assistance (Claude) was used for the implementation, validation and this description. The human submitter is responsible for reviewing and understanding every changed line.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (From [Docs] Add PR checklist skill for coding agents #57084, not yet merged.)

• AI assistance was used during the creation of this PR.

• Design Fit: Reuses ModernBertModel and the pooling API unchanged; the model is self-contained, with no core-component changes. The TP>1 limitation comes from existing ModernBERT code.

• Testing and Validation: Validated against the reference implementation on both public checkpoints and end to end. The added CI test is mutation-checked and has 6× tolerance margin.

• Code Quality and Style: Single-use helpers inlined; comments kept to non-obvious reasons; docs table, docs section and examples updated.

• Pull Request Contents: Summary, claims with evidence, reproduction commands, design tradeoffs, limitations, duplicate check and AI disclosure.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--58429.org.readthedocs.build/en/58429/