Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#1,343 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
5/5
Thời gian dự kiến
Hơn một tuần
Mức phù hợp với người mới
18/100
Loại issue
Tính năng
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
kotlin
Lĩnh vực
compilers, machine-learning

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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

Ngôn ngữ chính
Kotlin
Star
52
Fork
15
Merge trung bình
1 ngày 8 giờ
Pull request đã merge (30 ngày)
37

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của SKaiNET-developers/SKaiNET

Tất cả issue của SKaiNET-developers/SKaiNET

Issue tương tự

Thêm issue về Kotlin

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.