Make ModernBERT consume the validated shared attention layout and mask path

[forkwright]

Finding

The transformer crate implements head transposition, head merging, and attention-mask validation twice. The general Qwen/GQA path owns checked, Result-returning helpers and rejects a mask whose length differs from the sequence. The ModernBERT path carries a comment that its private helpers are “same as in attention.rs,” but its copies silently leave zero-filled output when a source range is missing and treat missing mask entries as padding.

One transformer representation therefore has two implementations with different failure semantics.

Verified against main e43777cd7f86674197d94e33533147dddfac4c82.

Evidence

• crates/transformers/src/attention.rs:128-165 validates x, mask.len() == seq, RoPE sequence capacity, and RoPE head width before attention proceeds.
• attention.rs:318-366 defines transpose_seq_heads and merge_heads as Result<Vec<f32>>; every source and destination slice is checked and a malformed range becomes Error::Shape.
• crates/transformers/src/modernbert.rs:345-369 validates only x.len() == seq * hidden. It does not validate mask.len() == seq before using the mask.
• modernbert.rs:381-393 splits Q/K/V with get(...).unwrap_or_default(), so an internal projection-length disagreement silently appends empty rows rather than producing the documented shape error.
• modernbert.rs:477-505 applies the mask with mask.get(j).copied().unwrap_or(0), converting every missing entry into padding instead of rejecting a malformed mask.
• modernbert.rs:510-543 labels its private head helpers “same as in attention.rs,” but returns plain Vec<f32> and copies only when both optional slices exist. A missing range leaves the pre-zeroed destination unchanged and the forward pass still returns Ok.
• Existing Match ModernBERT's layer-zero identity and zero-indexed global-attention schedule #7 and Load and execute the checkpoint's full ModernBERT prediction head #8 own ModernBERT architecture/checkpoint fidelity. mask_additive_in_place panics on division by zero for n == 0 or empty mask #31 and All-padding attention mask produces NaN softmax that propagates silently into encoder output #41 own separate low-level masking/softmax defects. No current issue owns this duplicated layout-and-validation path.

Why this matters

Shape and mask disagreements are model-integrity failures. Silent zero filling or implicit padding produces numerically valid tensors over the wrong inputs, which is more difficult to detect than a load/shape error. It can make the ModernBERT CPU path disagree with the shared transformer path while both appear successful, undermining its role as a correctness reference for later GPU work.

The copies also make every safety repair non-transitive. Tightening the checked helper or mask contract in attention.rs does not protect ModernBERT, as the current divergence already demonstrates.

Desired correction

Give the transformer crate one internal owner for:

• [seq, heads, d] ↔ [heads, seq, d] materialization;
• exact input/output length validation; and
• sequence-mask length validation.

Have both Qwen/GQA and ModernBERT consume that owner. Keep ModernBERT-specific local-window/global-attention masking as the legitimate delta, but require a validated [seq] mask before applying it.

Replace unwrap_or_default and optional-copy fallthrough in the forward path with typed Error::Shape propagation.

Done when:

• only one production implementation transposes and merges transformer heads;
• both attention families reject malformed tensor and mask lengths through the same error contract;
• ModernBERT cannot return Ok with zero-filled rows caused by missing slices;
• local/global-window policy remains explicit and separately tested; and
• adversarial fixtures prove a short mask and a malformed QKV/layout buffer fail loudly in both paths.