[Feature]: key-padding mask for MultiHeadAttention — batched encoder inference for BERT-class and Laya

Start with MultiHeadAttention.forward and HeadAttentionKernel, where the softmax and the causal/sliding-window masks are combined, then read BertEncoderRuntime.encode and LayaDecisionAdapter to see how one sequence is run per pass. Done means a null keyPaddingMask gives byte-identical output, BertEncoderRuntime.encodeBatch matches per-sequence output within 1e-5 on the LEAF goldens, and the six predict.json Laya cases pass as one batch.

🧠 D: DEFINE — Problem & Opportunity

transformer-core's MultiHeadAttention knows three mask shapes: causal, slidingWindow (with rightContext) and none. It has no key-padding mask, so a batch of sequences of different lengths cannot be attended correctly: padded key positions would receive attention weight. Every encoder consumer therefore runs one unpadded sequence per forward pass:

• BertEncoderRuntime.encode(tokenIds, attentionMask) takes one sequence (the mask is accepted but batching is not offered);
• LayaDecisionAdapter (#487) answers one question per encoder pass, batchSupport = false in its AdapterCapabilities; a request with twelve questions is twelve passes of a 322M-parameter encoder (measured 1.8–5 s each on a laptop, eager FP32).

Batching the questions of one request, or the items of one embedding sweep, needs the mask: pad to the longest sequence, mask the padded keys to -inf before the softmax, and ignore padded query rows in pooling / marker gathering.

🔍 A: ASSESS — Options

1. keyPaddingMask: Tensor? on MultiHeadAttention.forward (proposed): [batch, seqLen] of 0/1 (or lengths IntArray and the mask built inside), combined additively with the causal / sliding-window mask in HeadAttentionKernel. Zero cost when null.
2. Pack sequences without padding (block-diagonal attention): no wasted compute, but a different kernel and a larger change to the KV-cache path; the decoder runtimes do not need it.
3. Keep one sequence per pass and rely on the compiled / quantized lanes for speed: orthogonal, does not remove the N-passes-per-request problem.

✅ R: RESOLVE — Decision

Option 1, encoder path first (decoders keep their contract):

• MultiHeadAttention.forward(x, …, keyPaddingMask = null); HeadAttentionKernel adds -inf for masked keys before the softmax; works together with slidingWindow/rightContext (Laya's local layers) and causal = false.
• BertEncoderRuntime.encodeBatch(BatchEncoding) using the engine's BatchTokenizer.encodeBatch (padding to the longest) and masked mean pooling; LayaDecisionAdapter packs the questions of one request into one padded batch (batchSupport = true), marker gathering per row.
• Parity tests: batched output equals the per-sequence output within 1e-5 for BERT (LEAF goldens) and for Laya (predict.json cases run as one batch).

🧩 Lowering metadata — the mask as op metadata, not a baked tensor

Same principle as SKaiNET#1342 for GELU: the padding mask must survive direct execution, the recorded/optimized graph and every export with enough metadata that a lowering transformation can choose the legal form for its target. Where the engine stands today (checked 2026-10-14):

• TensorOps.scaledDotProductAttention(q, k, v, mask, scale, causal) is a single op. RecordingExecution records it as ScaledDotProductAttentionOperation with parameters = { scale, causal } and the explicit mask as a fourth operand; ComputeGraphExecutor dispatches on causal.
• transformer-core's MultiHeadAttention materialises its sliding/causal mask as an additive [1, 1, seqQ, seqKV] tensor (-1e30, buildSlidingCausalMask) and passes it as that operand (MultiHeadAttention.kt:483–489).
• The StableHLO converter (AttentionOperationsConverter) gives an explicit mask operand priority over its iota causal path, broadcasts it to the scores shape and adds it before the softmax. OpGranularityPolicy already names scaleddotproductattention as a fused-able op, but no converter calls keepFused yet.

What the padding mask should carry

A key-padding mask baked into the additive -1e30 tensor works on the eager path but throws the structure away: a compiled target sees an opaque [b, 1, Sq, Sk] constant-per-call, cannot recognise "these are padded keys", cannot fold it into a fused attention kernel that takes lengths, and the mask tensor itself becomes a graph input that scales with the batch. Proposal:

• Op parameters: parameters["maskKind"] = "none" | "causal" | "slidingWindow" | "keyPadding" | "combined", parameters["window"], parameters["rightContext"], and for padding the lengths operand ([batch] Int) instead of, or in addition to, the dense additive mask. The dense mask stays legal as the generic fallback (what MultiHeadAttention emits today), but the kind is recorded so an emitter can regenerate a cheaper form.
• One mask builder: the -1e30 additive construction moves into one builder (AttentionMask.additive(kind, lengths, window, rightContext, seqQ, seqKV)) that both the eager kernel (HeadAttentionKernel) and the StableHLO decomposition use. The converter's iota causal path and buildSlidingCausalMask are the same formula written twice today; the padding mask would be a third copy unless it is unified.

Emission rules (StableHLO), keyed on OpGranularityPolicy.keepFused("scaleddotproductattention")

A vendor compiler matches the composite name and attributes (lengths-based padding is what fused attention kernels take natively); everyone else inlines the decomposition or runs stablehlo-legalize-composite-to-call and gets the dense-mask formula back. The composite op is replaceable by its decomposition by definition, so semantics are identical on both rows.

Why this belongs in this issue

The acceptance list below adds keyPaddingMask to MultiHeadAttention and the kernel. If the first implementation bakes it into the dense additive tensor only, the compiled lane inherits a graph input of size b·Sq·Sk per layer and the fused-kernel path can never recover the lengths. Recording maskKind + lengths from the start costs one parameter and one operand; the dense fallback is unchanged. The transformers side of the work is therefore:

• MultiHeadAttention.forward(x, …, keyPaddingMask | keyLengths) records maskKind/lengths on the SDPA call (engine API addition: SKaiNET#1343 — keyLengths/window/rightContext on scaledDotProductAttention, one AttentionMask builder, lengths-based StableHLO lowering on the LoweringSpec/emitComposite seam of SKaiNET#1342).
• BertEncoderRuntime.encodeBatch and LayaDecisionAdapter pass lengths, not hand-built masks; masked pooling and marker gathering read the same lengths.
• Parity: eager dense-mask path == eager lengths path == compiled (DIRECT vs OPTIMIZED) within 1e-5 on the LEAF goldens and the six Laya predict.json cases.

🚀 C: COMMIT — Acceptance

• keyPaddingMask on MultiHeadAttention + kernel, null path byte-identical to today
• BertEncoderRuntime.encodeBatch, masked pooling, parity test on the LEAF goldens
• LayaDecisionAdapter batches questions per request, parity on the six predict.json cases, batchSupport = true
• Measured: time per question for a 6- and 12-question request before and after
• SDPA call records maskKind and a lengths operand (dense additive mask kept as fallback); one AttentionMask builder shared by the eager kernel and the StableHLO decomposition
• Compiled lane: padding mask emitted from lengths via iota/compare/select (no [b,1,Sq,Sk] graph input); composite form "skainet.sdpa" with mask_kind/window/right_context attributes behind OpGranularityPolicy
• Engine API issue for the lengths/maskKind parameters filed: SKaiNET#1343

Related: SKaiNET#1343 (engine side of this issue: SDPA lengths + mask metadata), SKaiNET#1342 (exact GELU op — the other engine gap the Laya adapter works around, and the home of the LoweringSpec/stablehlo.composite emission seam this note reuses), REQ-EXE notes in #487.