Skip to content

[Feature]: key-padding lengths for scaledDotProductAttention — mask as op metadata, lengths-based lowering #1343

Closed

Description

@michalharakal

🧠 D: DEFINE — Problem & Opportunity

TensorOps.scaledDotProductAttention(query, key, value, mask, scale, causal) knows two mask forms: the causal flag and an opaque additive mask tensor [batch, 1, seqLen, kvLen]. There is no notion of a key-padding mask, so a batch of sequences with different lengths can only be attended by baking -1e30 into that dense tensor on the caller's side.

That works on the eager path and loses the structure everywhere else:

  • the recorded graph (RecordingExecution → ScaledDotProductAttentionOperation, parameters = { scale, causal }) sees a per-call constant of size b·Sq·Sk per layer instead of "these keys are padding";
  • the StableHLO converter (AttentionOperationsConverter) can only broadcast and add it before the softmax; a fused attention kernel on a vendor target takes lengths, which can no longer be recovered;
  • transformer-core already carries its sliding/causal mask the same way (MultiHeadAttention.buildSlidingCausalMask, -1e30), and the converter has a second copy of the causal formula in its iota path — a padding mask would be the third copy.

Consumer: SKaiNET-transformers#490 (key-padding mask for MultiHeadAttention, batched BERT/Laya encoder inference) needs this to batch the questions of one typed-decision request instead of running one 322M-parameter encoder pass per question (1.8–5 s each, measured 2026-10-12/13).

🔍 A: ASSESS — Options

  1. keyLengths: Tensor<Int>? operand + recorded maskKind (proposed): the op gains a [batch] lengths operand and records maskKind ∈ { none, causal, slidingWindow, keyPadding, combined }, window, rightContext as parameters. The dense mask stays as the generic fallback, so no caller changes numerically.
  2. Dense mask only, as today: zero API change, but the compiled lane inherits a batch-sized graph input per layer and the fused-kernel path is closed for good.
  3. A separate maskedScaledDotProductAttention op: avoids widening the signature, but doubles the converter, the autograd entry and the backend kernels for one extra operand.

✅ R: RESOLVE — Decision

Option 1, with the mask construction unified:

  • TensorOps.scaledDotProductAttention(query, key, value, mask = null, scale = 0f, causal = false, keyLengths: Tensor<*, *>? = null, window: Int = 0, rightContext: Int = 0); VoidTensorOps, RecordingExecution, ComputeGraphExecutor, DefaultExecutionTape backward follow. RecordingExecution records maskKind, window, rightContext in parameters and keyLengths as operand 4.

  • One mask builder AttentionMask.additive(kind, lengths, window, rightContext, seqQ, seqKV) in skainet-lang-core, used by the CPU kernel (DefaultCpuOps.scaledDotProductAttention), by the StableHLO decomposition and, downstream, by MultiHeadAttention in place of buildSlidingCausalMask. Padded keys get -1e30 (the value the converter and transformer-core already agree on), never -inf, so a fully padded row stays finite.

  • CPU kernel: when keyLengths is given, the score loop stops at len[b] per batch row instead of adding a mask — cheaper than the dense path and bit-identical to it on the attended keys.

  • StableHLO lowering, keyed on OpGranularityPolicy.keepFused("scaleddotproductattention") (the policy is threaded into ConversionContext but unused so far — this and [Feature]: exact (erf) GELU as a first-class op — CPU kernels and HLO lowering #1342 are its first consumers):

    policy emitted form
    null / false (portable, llvm-cpu) decomposed as today; the padding mask is built from the keyLengths operand with iota + compare + select in the scores shape, 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 = { mask_kind = "keyPadding", causal = false, window = 128 : i32, right_context = 0 : i32, scale = 0.125 : f32 }, decomposition = @skainet_sdpa_…, version = 1 : i32 } with the decomposition function emitted by the same mask builder

    Both rows have identical semantics by construction: the StableHLO spec defines composite as replaceable by its decomposition, and stablehlo-legalize-composite-to-call does that mechanically. The LoweringSpec record and the ConversionContext.emitComposite helper are specified in [Feature]: exact (erf) GELU as a first-class op — CPU kernels and HLO lowering #1342; this issue reuses them for SDPA rather than adding a second mechanism.

  • Autograd: scaledDotProductAttentionBackward masks the gradient of padded keys to zero from the same builder.

  • Docs: operator matrix row, @Lowering(composite = "skainet.sdpa") once [Feature]: exact (erf) GELU as a first-class op — CPU kernels and HLO lowering #1342 adds the annotation.

🚀 C: COMMIT — Acceptance

  • keyLengths/window/rightContext on TensorOps.scaledDotProductAttention; dense mask path byte-identical to today
  • AttentionMask.additive builder shared by the CPU kernel, the StableHLO decomposition and the backward pass; the converter's iota causal path replaced by it
  • RecordingExecution records maskKind + lengths operand; ComputeGraphExecutor dispatches on them; DIRECT vs OPTIMIZED parity within 1e-6 on a padded batch
  • StableHLO: lengths-based decomposition and "skainet.sdpa" composite (version 1) behind OpGranularityPolicy; test runs the composite through stablehlo-legalize-composite-to-call and compares with the decomposed module
  • Parity test: padded batch of N sequences equals N unpadded single passes within 1e-5 (CPU, both kernel tiers)
  • transformers: MultiHeadAttention switches to the builder and passes lengths (SKaiNET-transformers#490)

Related: #1342 (exact GELU — same LoweringSpec / composite emission seam, first consumers of OpGranularityPolicy together), SKaiNET-transformers#490 (consumer), SKaiNET-transformers#487 (Laya adapter, batchSupport = false today).

Activity

  1. michalharakal commented on Oct 16, 2026

    @michalharakal
    ContributorAuthor

    Closing: both parts shipped in engine 0.60.0 and the consumer has landed.

    Left as noted in #1348: the SDPA backward rule stays a stub for the lengths path; the builder is in place for it.

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

    compute-backendCompute backends and graphsdsl/apiDSL and API designenhancementNew 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