Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

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

Aperta
  • Qualcosa non va?
#1,343 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
18/100
Tipo di issue
Funzionalità
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
kotlin
Ambito
compilers, machine-learning

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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

Lingua principale
Kotlin
Stelle
52
Fork
15
Merge medio
1g 4h
PR unite (30g)
37

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di SKaiNET-developers/SKaiNET

Tutte le issue di SKaiNET-developers/SKaiNET

Issue simili

Altre issue su Kotlin

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.