Skip to content

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

Closed

Description

@michalharakal

🧠 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 = 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")

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) 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.

Activity

  1. added
    enhancementNew feature or request
    codingImplementation task (DARC: C)
    skill:numericsPyTorch/NumPy/math background, no Kotlin required
    skill:kotlin-coreKotlin implementation in commonMain
    on Oct 13, 2026
  2. michalharakal commented on Oct 15, 2026

    @michalharakal
    ContributorAuthor

    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:

    request batched single
    9 mixed questions, one request 2472 ms/question 2735 ms/question
    3 billing questions 2138 ms/question 2355 ms/question
    2 structured-state questions 3013 ms/question 3190 ms/question
    2 CJK questions 2168 ms/question 2147 ms/question

    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 maskKind and 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.json 55 values, max |diff| 4.99e-5 (tolerance 1.5e-4); batched vs single logits 1.9e-6.

  3. michalharakal commented on Oct 16, 2026

    @michalharakal
    ContributorAuthor

    Closing: #496 is merged and goes out as 0.61.0 (#497), on engine 0.60.0.

    • PaddedBatch / PaddedBatchAware / withPaddedBatch in transformer-core; MultiHeadAttention runs a padded batch as one SDPA call with keyLengths, window, rightContext, so the recorded node carries maskKind and the compiled path takes the lengths as a graph input instead of a dense mask.
    • BertEncoderRuntime.forwardBatch / encodeBatch / encodePadded with masked pooling, DIRECT and OPTIMIZED.
    • LayaDecisionModel.*Batch, LayaAgent.logitsFor with length bucketing, LayaDecisionAdapter with batchSupport = true; the Laya fixtures reproduce through the batched path (55 values, max |diff| 4.99e-5).
    • Engine side (SKaiNET#1343, 0.60.0): AttentionMask rules shared by the CPU kernel and the StableHLO decomposition, lengths-based lowering, skainet.sdpa composite behind OpGranularityPolicy.

    Measured numbers are in the comment above: about 10 % per question on eager FP32 CPU, the gain is for the compiled lane and accelerators.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    codingImplementation task (DARC: C)enhancementNew feature or requestsize:m1-2 daysskill:kotlin-coreKotlin implementation in commonMainskill:numericsPyTorch/NumPy/math background, no Kotlin required

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions