Module model 

ModernBERT encoder and Laya decision heads, in candle.

Architecture follows upstream Laya and Hugging Face ModernBERT. Four details are easy to get wrong and are called out where they occur:

1. The encoder and scoring head use exact (erf) GELU, so gelu_erf — candle’s gelu is the tanh approximation.
2. The decision-head feed-forward uses ReLU, because upstream builds it from torch.nn.TransformerEncoderLayer, whose default activation is ReLU even though everything around it is GELU.
3. Padded queries are allowed to attend to valid keys. Without that a fully-masked softmax row produces NaN. Padded positions are never used as keys or pooled outputs, so valid-token results are unchanged.
4. RoPE tables are built in f32 and cast afterwards; building them directly in f16 loses precision at long positions.

Structs

AttentionMasks

Additive attention masks: 0 where attention is allowed, -inf elsewhere.

DecisionModel

The full Laya decision model: encoder, decision head, scorer and action head.

ModelOutput

One forward pass: masked scoring logits and action probabilities’ logits.

ModernBert

Rotary

Precomputed RoPE tables for one theta.