feat(ops)!: key-padding lengths and window metadata for scaledDotProductAttention (#1343, part 1)

CHANGELOG.md

@@ -2,6 +2,21 @@
 
 ## [Unreleased]
 
+### Added
+
+- **Key-padding lengths and window metadata for `scaledDotProductAttention`** (#1343, part 1): `scaledDotProductAttention(query, key, value, mask, scale, causal, keyLengths: Tensor<DType, *>? = null, window: Int = 0, rightContext: Int = 0)`. Padding (`keyLengths`, a `[batch]` integer tensor), the sliding-window band (`window`/`rightContext`, the `MultiHeadAttention` formula) and `causal` are *rules over positions* defined once in `AttentionMask` (`keyRange`, `attends`, `additive`, `MASKED = -1e30`), shared by the CPU kernel, the recorded graph and the StableHLO lowering. The CPU kernel computes scores only inside each row's attended key range — bit-identical to the dense `-1e30` mask on the attended keys, finite on a fully padded row — and a padded batch equals the unpadded single passes. The dense additive `mask` stays the generic fallback and is added on top of the rules.
+- **Mask rules as op metadata**: `RecordingExecution` and the KSP tracer record `parameters["maskKind" | "window" | "rightContext"]` (`maskKind` = `none | causal | slidingWindow | keyPadding | combined`, `AttentionMask.kindOf`) with `keyLengths` as the last operand after the optional dense mask; `ScaledDotProductAttentionOperation.operandLayout` resolves the layout for `ComputeGraphExecutor` and the converter. DIRECT vs OPTIMIZED parity on a padded batch within 1e-6.
+- **StableHLO lowering of the rules** (`SdpaLowering`, the single emitter-based builder behind `AttentionOperationsConverter`): `causal`, the window band and `keyLengths` become one `i1` keep mask from `stablehlo.iota`/`compare`/`and` and a `select 0 / -1e30` added to the scores, so no `[batch, 1, Sq, Sk]` constant or graph input enters the module; `keyLengths` is a plain `[batch] i32` input. The builder reads every shape from the operand MLIR types, which is what the `skainet.sdpa` composite decomposition function (part 2) needs. `TensorOps.scaledDotProductAttention` carries `@Lowering(composite = "skainet.sdpa")`.
+
+### Changed
+
+- `TensorOps.scaledDotProductAttention` gained `keyLengths`, `window` and `rightContext`: implementors of the interface must update the override (callers are source-compatible through the defaults). The generated `GraphDsl` (`skainet-lang-dag`) now types optional tensor operands (`mask`, `bias`, `keyLengths`) as `GraphValue<*>? = null` instead of mandatory positionals.
+- StableHLO SDPA export: the causal diagonal now carries the `Sk − Sq` offset the CPU kernel uses (a KV-cache decode step attends to its whole cache; `Sq == Sk` emission is unchanged); a dense mask and `causal = true` now apply together, as on the CPU, instead of the mask silently replacing the rule; positional/padding rules on a dynamic key length return a conversion failure instead of invalid `iota` IR (a dense mask still exports on the dynamic path).
+
+### Fixed
+
+- `DefaultCpuOps.scaledDotProductAttention` indexed a dense mask with `i % size`, which read batch row 0 for every batch: a `[batch, 1, Sq, Sk]` mask with `batch > 1` was silently wrong. The mask is now broadcast trailing-aligned (`[b | 1, H | nKV | 1, Sq | 1, Sk | 1]`, the StableHLO rule); the shapes that worked before index identically.
+
 ## [0.59.0] - 2026-10-14
 
 Headline: **exact GELU and lowering metadata (#1342).** `erf` becomes a first-class op and `gelu`

docs/modules/ROOT/pages/reference/operators/generated/index.adoc

@@ -1,6 +1,6 @@
 = AI-NET Operators Reference
 
-Generated from version `0.59.0` on 2026-10-14
+Generated from version `0.59.0` on 2026-10-15
 
 == Operators by Modality
 

docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc

@@ -1353,7 +1353,7 @@ fun cos(tensor:Tensor): Tensor
 
 [source,kotlin]
 ----
-fun scaledDotProductAttention(query:Tensor, key:Tensor, value:Tensor, mask:Tensor, scale:Float, causal:Boolean): Tensor
+fun scaledDotProductAttention(query:Tensor, key:Tensor, value:Tensor, mask:Tensor, scale:Float, causal:Boolean, keyLengths:Tensor, window:Int, rightContext:Int): Tensor
 ----
 
 === Parameters
@@ -1365,11 +1365,17 @@ fun scaledDotProductAttention(query:Tensor, key:Tensor, value:Tensor, mask:Tenso
 * `value: Tensor`
   [batch, nKVHeads, kvLen, headDim]
 * `mask: Tensor`
-  optional additive mask [batch, 1, seqLen, kvLen] (e.g. causal)
+  optional additive mask broadcastable to [batch, 1, seqLen, kvLen]
 * `scale: Float`
-  scaling factor, defaults to 1/sqrt(headDim)
+  scaling factor, `0` selects 1/sqrt(headDim)
 * `causal: Boolean`
-  if true, apply causal masking (ignore [mask] parameter)
+  if true, query `q` attends to keys `k <= kvLen - seqLen + q` only
+* `keyLengths: Tensor`
+  optional `[batch]` integer tensor: row `b` attends to keys `< keyLengths[b]`
+* `window: Int`
+  sliding-window size in keys, `0` for none: query `q` attends to keys `>= kvLen - seqLen + q - window + 1`
+* `rightContext: Int`
+  bounded lookahead of the window band in keys (needs `window > 0`): query `q` attends to keys `<= kvLen - seqLen + q + rightContext`
 
 === Return Type
 

docs/modules/ROOT/pages/reference/ops-status-matrix.adoc

@@ -1,7 +1,7 @@
 = Operator Coverage Matrix
 :description: Cross-backend status for every operator function in SKaiNET.
 
-Generated from `operators.json` version `0.59.0` on 2026-10-14.
+Generated from `operators.json` version `0.59.0` on 2026-10-15.
 
 Rows are `Operator.function` pairs. The `Validated` column shows whether the function's documentation has been DARC-validated by a reviewer (see xref:contributing/darc-workflow.adoc[DARC workflow]). The `Lowering` column lists the compile-path forms the function declares with `@Lowering` (`composite` = a `stablehlo.composite` when a target keeps it fused, `decompose` = StableHLO/CHLO primitives, `native` = a 1:1 target op); `—` means undeclared. Remaining columns are backends that appear in any function's `statusByBackend` map — a missing entry means the backend makes no claim about the function (treat it as "unknown", not "not supported").
 
@@ -73,10 +73,10 @@ Rows are `Operator.function` pairs. The `Validated` column shows whether the fun
 | `TensorOps.erf` | ✖ | composite · decompose | — | — | — 
 | `TensorOps.sin` | ✖ | — | — | — | — 
 | `TensorOps.cos` | ✖ | — | — | — | — 
-| `TensorOps.scaledDotProductAttention` | ✖ | — | — | — | — 
+| `TensorOps.scaledDotProductAttention` | ✖ | composite · decompose | — | — | — 
 | `Similarity.cosineDistance` | ✖ | — | ✅ | ✅ | ✅ 
 
-| *Done* | *1 / 66* | *2 / 66* | *1 / 66* | *1 / 66* | *1 / 66* 
+| *Done* | *1 / 66* | *3 / 66* | *1 / 66* | *1 / 66* | *1 / 66* 
 |===
 
 Per-function detail including notes lives in xref:reference/operators/generated/index.adoc[Operator reference].

skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api

@@ -339,7 +339,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/context/sched
 	public fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor;
 	protected final fun rewrapFeedOrderWeight (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 	public fun rsubScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
-	public fun scaledDotProductAttention (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;FZ)Lsk/ainet/lang/tensor/Tensor;
+	public fun scaledDotProductAttention (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;FZLsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor;
 	public fun sigmoid (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 	public fun sign (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 	public fun silu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;

skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt

@@ -3717,7 +3717,10 @@ public open class DefaultCpuOpsBase(
         value: Tensor<T, V>,
         mask: Tensor<T, V>?,
         scale: Float,
-        causal: Boolean
+        causal: Boolean,
+        keyLengths: Tensor<DType, *>?,
+        window: Int,
+        rightContext: Int
     ): Tensor<T, V> {
         // Expected shapes: [batch, heads, seqQ, headDim] for Q, [batch, heads, seqKV, headDim] for K/V
         require(query.rank == 4) { "SDPA: expected rank-4 query, got rank ${query.rank}" }
@@ -3742,12 +3745,30 @@ public open class DefaultCpuOpsBase(
         require(key.shape[2] == value.shape[2]) {
             "SDPA: K seqKV (${key.shape[2]}) does not match V seqKV (${value.shape[2]})"
         }
+        sk.ainet.lang.tensor.ops.AttentionMask.validate(window, rightContext)
 
         val batch = query.shape[0]
         val heads = query.shape[1]
         val seqQ = query.shape[2]
         val headDim = query.shape[3]
         val seqKV = key.shape[2]
+
+        // #1343: key padding and the window band are rules over positions, resolved per query row
+        // through AttentionMask.keyRange. Scores outside the row's range are never computed and
+        // get AttentionMask.MASKED, which is bit-identical to adding the dense -1e30 mask on the
+        // attended keys (the masked exp underflows to exactly 0 either way) and keeps a fully
+        // masked row finite. `causal` folds into the range on this path. Without lengths or a
+        // window the kernel below runs unchanged.
+        val lengths: IntArray? = keyLengths?.let { kl ->
+            require(kl.rank == 1 && kl.shape[0] == batch) { "SDPA: keyLengths must be [batch=$batch], got ${kl.shape}" }
+            val raw = kl.data.copyToFloatArray()
+            IntArray(batch) { b ->
+                val len = raw[b].toInt()
+                require(len in 0..seqKV) { "SDPA: keyLengths[$b] = $len is outside 0..$seqKV" }
+                len
+            }
+        }
+        val structured = lengths != null || window > 0
         // Grouped-query attention (SKEEP-005 phase 2): K/V are read in place through the
         // head-group index instead of being tiled to Q's head count upstream. Same loop
         // order and arithmetic as before, so nKVHeads == nHeads stays bit-identical.
@@ -3764,6 +3785,19 @@ public open class DefaultCpuOpsBase(
         val vBuf = value.data.copyToFloatArray()
         // Hoisted out of the per-head loop: the mask is read-only for the whole call.
         val maskBuf = mask?.data?.copyToFloatArray()
+        // Dense mask broadcast, trailing-aligned like the StableHLO lowering: [b | 1, H | nKV | 1, Sq | 1, Sk | 1].
+        // (The former `i % size` read batch row 0 for every batch, so a [batch, 1, Sq, Sk] mask with batch > 1
+        // was silently wrong; shapes that worked — [1, 1, Sq, Sk], [1, 1, 1, Sk], [Sq, Sk] — index identically here.)
+        val maskDims = mask?.let { m ->
+            require(m.rank in 1..4) { "SDPA: mask must have rank 1..4, got ${m.rank}" }
+            val d = IntArray(4) { 1 }
+            for (i in 0 until m.rank) d[4 - m.rank + i] = m.shape[i]
+            require(d[0] == 1 || d[0] == batch) { "SDPA: mask batch (${d[0]}) must be 1 or $batch" }
+            require(d[1] == 1 || d[1] == heads || d[1] == key.shape[1]) { "SDPA: mask heads (${d[1]}) must be 1, ${key.shape[1]} or $heads" }
+            require(d[2] == 1 || d[2] == seqQ) { "SDPA: mask seqQ (${d[2]}) must be 1 or $seqQ" }
+            require(d[3] == 1 || d[3] == seqKV) { "SDPA: mask seqKV (${d[3]}) must be 1 or $seqKV" }
+            d
+        }
 
         val outBuf = FloatArray(batch * heads * seqQ * headDim)
 
@@ -3784,21 +3818,41 @@ public open class DefaultCpuOpsBase(
                 val b = unit / heads
                 val h = unit % heads
                 val kvH = h / nRep
-                // Compute attention scores: Q @ K^T, then scale
-                for (qi in 0 until seqQ) {
-                    for (ki in 0 until seqKV) {
-                        var dot = 0f
+                if (structured) {
+                    // Scores only inside the row's attended key range; the rest is MASKED.
+                    val length = lengths?.get(b) ?: seqKV
+                    for (qi in 0 until seqQ) {
+                        val range = sk.ainet.lang.tensor.ops.AttentionMask.keyRange(qi, seqQ, seqKV, causal, window, rightContext, length)
+                        val row = qi * seqKV
+                        scores.fill(sk.ainet.lang.tensor.ops.AttentionMask.MASKED, row, row + seqKV)
                         val qOff = ((b * heads + h) * seqQ + qi) * headDim
-                        val kOff = ((b * kvHeads + kvH) * seqKV + ki) * headDim
-                        for (d in 0 until headDim) {
-                            dot += qBuf[qOff + d] * kBuf[kOff + d]
+                        for (ki in range) {
+                            var dot = 0f
+                            val kOff = ((b * kvHeads + kvH) * seqKV + ki) * headDim
+                            for (d in 0 until headDim) {
+                                dot += qBuf[qOff + d] * kBuf[kOff + d]
+                            }
+                            scores[row + ki] = dot * effectiveScale
+                        }
+                    }
+                } else {
+                    // Compute attention scores: Q @ K^T, then scale
+                    for (qi in 0 until seqQ) {
+                        for (ki in 0 until seqKV) {
+                            var dot = 0f
+                            val qOff = ((b * heads + h) * seqQ + qi) * headDim
+                            val kOff = ((b * kvHeads + kvH) * seqKV + ki) * headDim
+                            for (d in 0 until headDim) {
+                                dot += qBuf[qOff + d] * kBuf[kOff + d]
+                            }
+                            scores[qi * seqKV + ki] = dot * effectiveScale
                         }
-                        scores[qi * seqKV + ki] = dot * effectiveScale
                     }
                 }
 
-                // Apply causal mask (set future positions to -inf)
-                if (causal) {
+                // Apply causal mask (set future positions to -inf); folded into the key range above
+                // when the structured path is active.
+                if (causal && !structured) {
                     for (qi in 0 until seqQ) {
                         // For single-token inference with KV cache: qi indexes from the
                         // query's perspective. The valid range of ki is [0, seqKV).
@@ -3811,10 +3865,21 @@ public open class DefaultCpuOpsBase(
                     }
                 }
 
-                // Apply external mask if provided
-                if (maskBuf != null) {
-                    for (i in scores.indices) {
-                        scores[i] += maskBuf[i % maskBuf.size]
+                // Apply external mask if provided (broadcast per maskDims)
+                if (maskBuf != null && maskDims != null) {
+                    val mb = maskDims[0]; val mh = maskDims[1]; val mq = maskDims[2]; val mk = maskDims[3]
+                    val bI = if (mb == 1) 0 else b
+                    val hI = when (mh) { 1 -> 0; heads -> h; else -> kvH }
+                    val base = (bI * mh + hI) * mq
+                    for (qi in 0 until seqQ) {
+                        val rowM = (base + (if (mq == 1) 0 else qi)) * mk
+                        val row = qi * seqKV
+                        if (mk == 1) {
+                            val mv = maskBuf[rowM]
+                            for (ki in 0 until seqKV) scores[row + ki] += mv
+                        } else {
+                            for (ki in 0 until seqKV) scores[row + ki] += maskBuf[rowM + ki]
+                        }
                     }
                 }