Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

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

Open
  • Is something wrong?
#1,343 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
18/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
kotlin
Domain
compilers, machine-learning

Research direction

Start from TensorOps.scaledDotProductAttention, then follow the mask handling in DefaultCpuOps.scaledDotProductAttention, RecordingExecution, and AttentionOperationsConverter, which the issue names. The issue depends on the not-yet-built LoweringSpec from #1342, so check that first. Done means the acceptance checklist holds: a padded batch matches unpadded single passes within 1e-5, and the dense mask path is unchanged.

Written by the indexing model from the issue text.

Description

compute-backend dsl/api enhancement size:m skill:kotlin-core skill:numerics

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

Dominant language
Kotlin
Stars
52
Forks
15
Avg merge
1d 4h
Merged PRs (30d)
37

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from SKaiNET-developers/SKaiNET

All issues in SKaiNET-developers/SKaiNET

Similar issues

More Kotlin issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.