[Feature]: key-padding lengths for scaledDotProductAttention — mask as op metadata, lengths-based lowering
メンテナーはふだん 1 日以内に返信
まだ誰も着手していません。
評価
- 難易度
- 5/5
- 見積もり時間
- 1週間以上
- 初心者へのやさしさ
- 18/100
- issue の種類
- 機能追加
- 明瞭さ
- おおむね明確
- 活発さ
- 活発
- 技術スタック
- kotlin
調査の方向性
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.
索引モデルが issue の本文から書いたものです。
説明
🧠 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 sizeb·Sq·Skper 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-corealready 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
keyLengths: Tensor<Int>?operand + recordedmaskKind(proposed): the op gains a[batch]lengths operand and recordsmaskKind ∈ { none, causal, slidingWindow, keyPadding, combined },window,rightContextas parameters. The densemaskstays as the generic fallback, so no caller changes numerically.- 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.
- A separate
maskedScaledDotProductAttentionop: 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,DefaultExecutionTapebackward follow.RecordingExecutionrecordsmaskKind,window,rightContextinparametersandkeyLengthsas operand 4. -
One mask builder
AttentionMask.additive(kind, lengths, window, rightContext, seqQ, seqKV)inskainet-lang-core, used by the CPU kernel (DefaultCpuOps.scaledDotProductAttention), by the StableHLO decomposition and, downstream, byMultiHeadAttentionin place ofbuildSlidingCausalMask. 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
keyLengthsis given, the score loop stops atlen[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 intoConversionContextbut 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 keyLengthsoperand withiota+compare+selectin the scores shape, so no[b, 1, Sq, Sk]constant enters the moduletrue(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 builderBoth rows have identical semantics by construction: the StableHLO spec defines
compositeas replaceable by itsdecomposition, andstablehlo-legalize-composite-to-calldoes that mechanically. TheLoweringSpecrecord and theConversionContext.emitCompositehelper are specified in #1342; this issue reuses them for SDPA rather than adding a second mechanism. -
Autograd:
scaledDotProductAttentionBackwardmasks 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/rightContextonTensorOps.scaledDotProductAttention; densemaskpath byte-identical to today -
AttentionMask.additivebuilder shared by the CPU kernel, the StableHLO decomposition and the backward pass; the converter's iota causal path replaced by it -
RecordingExecutionrecordsmaskKind+ lengths operand;ComputeGraphExecutordispatches on them; DIRECT vs OPTIMIZED parity within 1e-6 on a padded batch - StableHLO: lengths-based decomposition and
"skainet.sdpa"composite (version 1) behindOpGranularityPolicy; test runs the composite throughstablehlo-legalize-composite-to-calland 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:
MultiHeadAttentionswitches 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).
- 主要言語
- Kotlin
- スター
- 52
- フォーク
- 15
- 平均マージ
- 1日 4時間
- マージ済み PR(30日)
- 37
環境構築
- Dockerfile・Docker Compose ファイルなし
- プルリクエストのテンプレートなし
- コントリビューションガイドを読む
はじめの一歩
- issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
- 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
- リポジトリをフォークし、ブランチを切って変更します。
- issue 番号を参照したプルリクエストを送ります。
SKaiNET-developers/SKaiNET のほかの issue
-
coding good first issue platform size:xs skill:js sub-issue
難易度 2/5 1〜3時間 初心者へのやさしさ 72/100
SKaiNET-developers/SKaiNET#1232 ·
メンテナーはふだん 1 日以内に返信
-
skeep tracking
難易度 5/5 1週間以上 初心者へのやさしさ 15/100
SKaiNET-developers/SKaiNET#1331 · コメント 1 件 ·
メンテナーはふだん 1 日以内に返信
-
assessment size:s skill:review sub-issue
難易度 4/5 1〜2日 初心者へのやさしさ 18/100
SKaiNET-developers/SKaiNET#1330 ·
メンテナーはふだん 1 日以内に返信
-
documentation good first issue size:s skill:docs sub-issue
難易度 3/5 1〜2日 初心者へのやさしさ 78/100
SKaiNET-developers/SKaiNET#1329 ·
メンテナーはふだん 1 日以内に返信
-
coding good first issue size:s skill:kotlin-core sub-issue
難易度 3/5 半日 初心者へのやさしさ 75/100
SKaiNET-developers/SKaiNET#1328 ·
メンテナーはふだん 1 日以内に返信
SKaiNET-developers/SKaiNET の issue をすべて見る
似ている issue
-
Dev
難易度 2/5 1〜3時間 初心者へのやさしさ 62/100
ankidroid/Anki-Android#22555 ·
メンテナーはふだん 1 日以内に返信
-
enhancement good first issue help wanted
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
-
core: RAC_LLM_OPTIONS_DEFAULT.min_p disagrees with the IDL's annotated default対応中かも このイシューにリンクされたプルリクエストがオープン中、またはマージ済みです。 オープンbug core p3
難易度 2/5 1〜3時間 初心者へのやさしさ 70/100
RunanywhereAI/runanywhere-sdks#1083 ·
メンテナーはふだん 5 日以内に返信
-
enhancement
難易度 2/5 1〜3時間 初心者へのやさしさ 75/100
メンテナーはふだん 1 日以内に返信
-
Java and Kotlin (JVM): parse gradle/wrapper/gradle-wrapper.properties with exact line locationsオープンarea: detector eco: java help wanted lane: free-agent roadmap size: S type: parser
難易度 2/5 1〜3時間 初心者へのやさしさ 62/100
arminazimi/evergreen#355 ·
メンテナーはふだん 1 日以内に返信