Repository navigation
[Feature]: key-padding mask for MultiHeadAttention — batched encoder inference for BERT-class and Laya #490
Description
🧠 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(Typed decisions (choice / score / boolean probability) and the Laya adapter #487) answers one question per encoder pass,batchSupport = falsein itsAdapterCapabilities; 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
keyPaddingMask: Tensor?onMultiHeadAttention.forward(proposed):[batch, seqLen]of 0/1 (or lengthsIntArrayand the mask built inside), combined additively with the causal / sliding-window mask inHeadAttentionKernel. Zero cost when null.- 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.
- 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);HeadAttentionKerneladds-inffor masked keys before the softmax; works together withslidingWindow/rightContext(Laya's local layers) andcausal = false.BertEncoderRuntime.encodeBatch(BatchEncoding)using the engine'sBatchTokenizer.encodeBatch(padding to the longest) and masked mean pooling;LayaDecisionAdapterpacks 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.jsoncases 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.RecordingExecutionrecords it asScaledDotProductAttentionOperationwithparameters = { scale, causal }and the explicitmaskas a fourth operand;ComputeGraphExecutordispatches oncausal.transformer-core'sMultiHeadAttentionmaterialises 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.OpGranularityPolicyalready namesscaleddotproductattentionas a fused-able op, but no converter callskeepFusedyet.
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 (whatMultiHeadAttentionemits today), but the kind is recorded so an emitter can regenerate a cheaper form. - One mask builder: the
-1e30additive 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 andbuildSlidingCausalMaskare 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")
| policy | emitted form |
|---|---|
null / false (portable, llvm-cpu) |
decomposed as today: dot_general, additive mask, softmax, dot_general. The padding mask is built from the lengths operand with iota + compare + select (same shape as the existing iota causal path), so no [b, 1, Sq, Sk] constant enters the module |
true (vendor target with a fused attention kernel) |
stablehlo.composite "skainet.sdpa" %q, %k, %v, %lengths { composite_attributes = { causal = false, window = 128 : i32, right_context = 0 : i32, mask_kind = "keyPadding", scale = 0.125 : f32 }, decomposition = @skainet_sdpa_…, version = 1 : i32 } with the decomposition function built by the same mask builder |
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)recordsmaskKind/lengthson the SDPA call (engine API addition: SKaiNET#1343 —keyLengths/window/rightContextonscaledDotProductAttention, oneAttentionMaskbuilder, lengths-based StableHLO lowering on theLoweringSpec/emitCompositeseam of SKaiNET#1342).BertEncoderRuntime.encodeBatchandLayaDecisionAdapterpass 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.jsoncases.
🚀 C: COMMIT — Acceptance
-
keyPaddingMaskonMultiHeadAttention+ kernel, null path byte-identical to today -
BertEncoderRuntime.encodeBatch, masked pooling, parity test on the LEAF goldens -
LayaDecisionAdapterbatches questions per request, parity on the sixpredict.jsoncases,batchSupport = true - Measured: time per question for a 6- and 12-question request before and after
- SDPA call records
maskKindand alengthsoperand (dense additive mask kept as fallback); oneAttentionMaskbuilder shared by the eager kernel and the StableHLO decomposition - Compiled lane: padding mask emitted from
lengthsvia iota/compare/select (no[b,1,Sq,Sk]graph input); composite form"skainet.sdpa"withmask_kind/window/right_contextattributes behindOpGranularityPolicy - Engine API issue for the
lengths/maskKindparameters 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.
Activity
michalharakal commentedon Oct 15, 2026
Measured numbers for the time per question acceptance item, eager FP32 CPU on a laptop, Laya multilingual checkpoint, through the padded-batch path of #496:
Batching saves about 10 % per question at best, and padding to the longest sequence without length bucketing was 2.4× slower than N single passes. The eager path is compute-bound: the cost of N passes is the FLOPs of N sequences, not per-pass overhead, so the "twelve passes of a 322M encoder" framing in the DEFINE section overstates what padding alone can recover on this lane. The adapter therefore buckets by length (
count · longest ≤ maxBatchTokens, padding ≤ 25 % of the real tokens, lone sequences keep the fused single path).What the batch does buy, and why the lengths-as-metadata design still matters: one SDPA node per layer with
maskKindand a lengths operand, the compiled graph takes the lengths as an input instead of a dense[b, 1, S, S]mask, and fused attention kernels on accelerators consume exactly that form. The speed case for #490 is on the compiled lane and on accelerators, not on eager CPU.Parity through the batched path:
predict.json55 values, max |diff| 4.99e-5 (tolerance 1.5e-4); batched vs single logits 1.9e-6.feat(encoders): padded batches with key lengths for BERT and Laya (#490)
michalharakal commentedon Oct 16, 2026
Closing: #496 is merged and goes out as 0.61.0 (#497), on engine 0.60.0.
PaddedBatch/PaddedBatchAware/withPaddedBatchintransformer-core;MultiHeadAttentionruns a padded batch as one SDPA call withkeyLengths,window,rightContext, so the recorded node carriesmaskKindand the compiled path takes the lengths as a graph input instead of a dense mask.BertEncoderRuntime.forwardBatch/encodeBatch/encodePaddedwith masked pooling, DIRECT and OPTIMIZED.LayaDecisionModel.*Batch,LayaAgent.logitsForwith length bucketing,LayaDecisionAdapterwithbatchSupport = true; the Laya fixtures reproduce through the batched path (55 values, max |diff| 4.99e-5).AttentionMaskrules shared by the CPU kernel and the StableHLO decomposition, lengths-based lowering,skainet.sdpacomposite behindOpGranularityPolicy.Measured numbers are in the comment above: about 10 % per question on eager FP32 CPU, the gain is for the compiled lane and accelerators.