From b17ddca394e535f467a44bc2d8959872cae72d33 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Thu, 15 Oct 2026 17:15:17 +0200 Subject: [PATCH 1/2] feat(ops)!: key-padding lengths and window metadata for scaledDotProductAttention (#1343, part 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Padding (`keyLengths`, a [batch] integer tensor), the sliding-window band (`window`/`rightContext`) and `causal` become rules over positions defined once in `AttentionMask` and shared by the CPU kernel, the recorded graph and the StableHLO lowering. The dense additive mask stays the generic fallback. - `TensorOps.scaledDotProductAttention(…, keyLengths = null, window = 0, rightContext = 0)`, `@Lowering(composite = "skainet.sdpa")`; VoidTensorOps, RecordingExecution, ComputeGraphExecutor, the KSP tracer (`OpAttributeFactory.sdpaMaskKind`) and the graph DSL follow. - Recorded metadata: `maskKind` (none | causal | slidingWindow | keyPadding | combined), `window`, `rightContext`; `keyLengths` is the last operand after the optional dense mask (`ScaledDotProductAttentionOperation.operandLayout`). - CPU: scores only inside each row's attended key range, MASKED (-1e30) elsewhere — bit-identical to the dense mask on the attended keys, finite on a fully padded row. Fixed the dense-mask indexing, which read batch row 0 for every batch; masks now broadcast trailing-aligned. - StableHLO: converter body moved into the emitter-based `SdpaLowering.emit`, which reads shapes from the operand MLIR types; rules lower to one i1 keep mask (iota/compare/and) and select 0 / -1e30. The causal diagonal carries the Sk − Sq offset the CPU uses; a dense mask and `causal` apply together; rules on a dynamic key length fail the conversion instead of emitting invalid iota IR. - Graph DSL: optional tensor operands (mask, bias, keyLengths) are `GraphValue<*>? = null` instead of mandatory positionals. - Tests: AttentionMaskTest, SdpaKeyLengthsTest, SdpaKeyLengthsRecordingTest, SdpaMaskRulesHloTest; api dumps, generated docs, changelog. --- CHANGELOG.md | 15 + .../reference/operators/generated/index.adoc | 2 +- .../operators/generated/tensorops.adoc | 14 +- .../pages/reference/ops-status-matrix.adoc | 6 +- .../api/jvm/skainet-backend-cpu.api | 2 +- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 95 ++++- .../exec/tensor/ops/SdpaKeyLengthsTest.kt | 146 ++++++++ .../sk/ainet/tape/RecordingExecution.kt | 15 +- .../lang/graph/exec/ComputeGraphExecutor.kt | 11 +- .../compile/graph/ComputeGraphExecutorTest.kt | 2 +- .../ainet/exec/SdpaKeyLengthsRecordingTest.kt | 110 ++++++ .../api/jvm/skainet-compile-hlo.api | 8 + .../AttentionOperationsConverter.kt | 273 +++------------ .../compile/hlo/converters/SdpaLowering.kt | 326 ++++++++++++++++++ .../ainet/compile/hlo/SdpaMaskRulesHloTest.kt | 102 ++++++ .../api/jvm/skainet-lang-core.api | 49 ++- .../sk/ainet/lang/tensor/ops/AttentionMask.kt | 104 ++++++ .../ainet/lang/tensor/ops/TensorOperations.kt | 33 ++ .../sk/ainet/lang/tensor/ops/TensorOps.kt | 27 +- .../sk/ainet/lang/tensor/ops/VoidTensorOps.kt | 5 +- .../sk/ainet/lang/trace/OpAttributeFactory.kt | 8 + .../lang/tensor/ops/AttentionMaskTest.kt | 77 +++++ .../api/jvm/skainet-lang-dag.api | 4 +- .../sk/ainet/lang/ksp/GraphDslGenerator.kt | 6 +- .../sk/ainet/lang/ksp/MethodAnalyzer.kt | 6 + .../ainet/lang/ksp/TracingWrapperProcessor.kt | 2 +- .../sk/ainet/lang/trace/OpAttributeFactory.kt | 3 + 27 files changed, 1181 insertions(+), 270 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SdpaKeyLengthsTest.kt create mode 100644 skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/SdpaKeyLengthsRecordingTest.kt create mode 100644 skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt create mode 100644 skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/SdpaMaskRulesHloTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/AttentionMask.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 9430de31d..147121745 100644 --- a/CHANGELOG.md +++ b/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? = 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` diff --git a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc index d7f0c8256..d574034e7 100644 --- a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc +++ b/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 diff --git a/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc b/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc index 2f5241233..38c9e227c 100644 --- a/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc +++ b/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 diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc index 56368ed48..6dd838fca 100644 --- a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc +++ b/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]. diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 0ea2ba696..9fc00dfb8 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/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; diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 6ff2f1625..9555039d9 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/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, mask: Tensor?, scale: Float, - causal: Boolean + causal: Boolean, + keyLengths: Tensor?, + window: Int, + rightContext: Int ): Tensor { // 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] + } } } diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SdpaKeyLengthsTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SdpaKeyLengthsTest.kt new file mode 100644 index 000000000..e1dac5fca --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SdpaKeyLengthsTest.kt @@ -0,0 +1,146 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.AttentionMask +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1343: key padding, the window band and `causal` as rules over positions on the CPU kernel. + * The structured path must equal the dense `-1e30` mask the caller used to build by hand, bit + * for bit, and a padded batch must equal the unpadded single passes. + */ +class SdpaKeyLengthsTest { + + private fun fill(size: Int, seed: Int) = FloatArray(size) { i -> (((i * 31 + seed * 17) % 23) - 11) / 11f } + + private class Inputs(val q: Tensor, val k: Tensor, val v: Tensor) + + private fun inputs(ctx: ExecutionContext, batch: Int, heads: Int, kvHeads: Int, seqQ: Int, seqKV: Int, headDim: Int, seed: Int = 0) = Inputs( + ctx.fromFloatArray(Shape(batch, heads, seqQ, headDim), FP32::class, fill(batch * heads * seqQ * headDim, 1 + seed)), + ctx.fromFloatArray(Shape(batch, kvHeads, seqKV, headDim), FP32::class, fill(batch * kvHeads * seqKV * headDim, 2 + seed)), + ctx.fromFloatArray(Shape(batch, kvHeads, seqKV, headDim), FP32::class, fill(batch * kvHeads * seqKV * headDim, 3 + seed)), + ) + + @Suppress("UNCHECKED_CAST") + private fun lengthsTensor(ctx: ExecutionContext, lengths: IntArray): Tensor = + ctx.fromIntArray(Shape(lengths.size), Int32::class, lengths) as Tensor + + private fun denseMask(ctx: ExecutionContext, batch: Int, seqQ: Int, seqKV: Int, causal: Boolean, window: Int, rightContext: Int, lengths: IntArray?): Tensor = + ctx.fromFloatArray(Shape(batch, 1, seqQ, seqKV), FP32::class, AttentionMask.additive(batch, seqQ, seqKV, causal, window, rightContext, lengths)) + + private fun floats(t: Tensor<*, *>) = t.data.copyToFloatArray() + + @Test + fun lengths_path_equals_the_dense_mask_bit_for_bit() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + for ((heads, kvHeads) in listOf(4 to 4, 4 to 2)) for ((seqQ, seqKV) in listOf(6 to 6, 3 to 10, 1 to 7)) for (causal in listOf(false, true)) { + val batch = 3 + val lengths = intArrayOf(seqKV, (seqKV + 1) / 2, 1) + val x = inputs(ctx, batch, heads, kvHeads, seqQ, seqKV, 8) + val structured = ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, mask = null, scale = 0f, causal = causal, keyLengths = lengthsTensor(ctx, lengths)) + val dense = ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, mask = denseMask(ctx, batch, seqQ, seqKV, causal, 0, 0, lengths), scale = 0f, causal = false) + assertContentEquals(floats(dense), floats(structured), "heads=$heads kv=$kvHeads seqQ=$seqQ seqKV=$seqKV causal=$causal") + } + } + + @Test + fun window_band_equals_the_dense_mask_bit_for_bit() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + for ((seqQ, seqKV) in listOf(8 to 8, 1 to 16, 5 to 12)) for ((window, rightContext) in listOf(3 to 0, 4 to 2, 16 to 4)) { + val x = inputs(ctx, 2, 4, 2, seqQ, seqKV, 8) + val structured = ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, window = window, rightContext = rightContext) + val dense = ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, mask = denseMask(ctx, 2, seqQ, seqKV, false, window, rightContext, null)) + assertContentEquals(floats(dense), floats(structured), "seqQ=$seqQ seqKV=$seqKV window=$window rightContext=$rightContext") + // Window + padding together ("combined"). + val lengths = intArrayOf(seqKV, seqKV / 2) + val both = ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, keyLengths = lengthsTensor(ctx, lengths), window = window, rightContext = rightContext) + val bothDense = ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, mask = denseMask(ctx, 2, seqQ, seqKV, false, window, rightContext, lengths)) + assertContentEquals(floats(bothDense), floats(both), "combined seqQ=$seqQ seqKV=$seqKV window=$window") + } + } + + @Test + fun padded_batch_equals_the_unpadded_single_passes() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + val heads = 4; val headDim = 8; val maxLen = 10 + val lengths = intArrayOf(10, 4, 7) + // Build each sequence alone, then the padded batch from the same rows (padding rows are junk). + val singles = lengths.indices.map { b -> inputs(ctx, 1, heads, heads, lengths[b], lengths[b], headDim, seed = b * 10) } + fun pad(parts: List>, seqIdx: Int): Tensor { + val out = FloatArray(lengths.size * heads * maxLen * headDim) { 7f } + for (b in lengths.indices) { + val src = floats(parts[b]); val len = lengths[b] + for (h in 0 until heads) for (s in 0 until len) for (d in 0 until headDim) + out[((b * heads + h) * maxLen + s) * headDim + d] = src[((h * len) + s) * headDim + d] + } + return ctx.fromFloatArray(Shape(lengths.size, heads, maxLen, headDim), FP32::class, out) + } + val q = pad(singles.map { it.q }, 2); val k = pad(singles.map { it.k }, 2); val v = pad(singles.map { it.v }, 2) + for (causal in listOf(false, true)) { + val batched = floats(ctx.ops.scaledDotProductAttention(q, k, v, causal = causal, keyLengths = lengthsTensor(ctx, lengths))) + for (b in lengths.indices) { + val single = floats(ctx.ops.scaledDotProductAttention(singles[b].q, singles[b].k, singles[b].v, causal = causal)) + val len = lengths[b] + var maxDiff = 0f + for (h in 0 until heads) for (s in 0 until len) for (d in 0 until headDim) { + val got = batched[((b * heads + h) * maxLen + s) * headDim + d] + val want = single[((h * len) + s) * headDim + d] + maxDiff = maxOf(maxDiff, kotlin.math.abs(got - want)) + } + assertTrue(maxDiff <= 1e-6f, "row $b causal=$causal max |Δ| = $maxDiff") + } + } + } + + @Test + fun fully_padded_row_stays_finite() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + val x = inputs(ctx, 2, 2, 2, 3, 5, 4) + val out = floats(ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, keyLengths = lengthsTensor(ctx, intArrayOf(0, 5)))) + assertTrue(out.all { it.isFinite() }, "fully padded row produced a non-finite value") + // Uniform attention over the keys, as the dense -1e30 mask yields. + val dense = floats(ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, mask = denseMask(ctx, 2, 3, 5, false, 0, 0, intArrayOf(0, 5)))) + assertContentEquals(dense, out) + } + + @Test + fun structured_path_is_schedule_independent() { + val x0 = inputs(DirectCpuExecutionContext(schedule = Schedule.Sequential), 2, 8, 2, 64, 64, 32) + val seq = DirectCpuExecutionContext(schedule = Schedule.Sequential) + val par = DirectCpuExecutionContext(schedule = object : Schedule { + override val parallelism = 3 + override val name = "reversing" + override fun forRange(n: Int, grain: Int, body: (Int, Int) -> Unit) { + val tasks = Schedule.tasksFor(n, grain, parallelism); if (tasks == 0) return + val chunk = Schedule.chunkFor(n, tasks) + (0 until n step chunk).map { s -> s to minOf(s + chunk, n) }.reversed().forEach { (s, e) -> body(s, e) } + } + }) + val lengths = intArrayOf(64, 17) + val a = floats(seq.ops.scaledDotProductAttention(x0.q, x0.k, x0.v, causal = true, keyLengths = lengthsTensor(seq, lengths), window = 16)) + val b = floats(par.ops.scaledDotProductAttention(x0.q, x0.k, x0.v, causal = true, keyLengths = lengthsTensor(par, lengths), window = 16)) + assertContentEquals(a, b) + } + + @Test + fun rejects_bad_lengths_and_window_arguments() { + val ctx = DirectCpuExecutionContext() + val x = inputs(ctx, 2, 2, 2, 3, 5, 4) + assertFailsWith { ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, keyLengths = lengthsTensor(ctx, intArrayOf(5))) } + assertFailsWith { ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, keyLengths = lengthsTensor(ctx, intArrayOf(5, 6))) } + assertFailsWith { ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, window = -1) } + assertFailsWith { ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, rightContext = 1) } + assertEquals(Shape(2, 2, 3, 4), ctx.ops.scaledDotProductAttention(x.q, x.k, x.v, window = 2, rightContext = 1).shape) + } +} diff --git a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt index df1912047..a735ae5b5 100644 --- a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt +++ b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt @@ -486,15 +486,16 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor ): Tensor = base.convTranspose1d(input, weight, bias, stride, padding, outputPadding, dilation, groups) override fun scaledDotProductAttention( query: Tensor, key: Tensor, value: Tensor, - mask: Tensor?, scale: Float, causal: Boolean + mask: Tensor?, scale: Float, causal: Boolean, + keyLengths: Tensor?, window: Int, rightContext: Int ): Tensor { - val out = base.scaledDotProductAttention(query, key, value, mask, scale, causal) - val params = mutableMapOf( - "scale" to scale, - "causal" to causal - ) + val out = base.scaledDotProductAttention(query, key, value, mask, scale, causal, keyLengths, window, rightContext) + // #1343: the mask rules travel as parameters (`maskKind` derived from them) and `keyLengths` + // as the last operand, after the optional dense mask — the layout + // ScaledDotProductAttentionOperation.operandLayout resolves. + val params = ScaledDotProductAttentionOperation.parameters(scale, causal, window, rightContext, keyLengths != null) @Suppress("UNCHECKED_CAST") - val inputs = listOfNotNull(query, key, value, mask) as List> + val inputs = listOfNotNull(query, key, value, mask, keyLengths) as List> record(ScaledDotProductAttentionOperation(params), inputs, listOf(out)) return out } diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt index b07902ece..ac68b1c46 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt @@ -265,13 +265,20 @@ public class ComputeGraphExecutor( "scaled_dot_product_attention", "scaledDotProductAttention", "sdpa" -> { val causal = params["causal"] as? Boolean ?: false val scale = params["scale"] as? Float + // #1343: `keyLengths` is the last operand when `maskKind` says so, the dense mask operand 3. + val layout = sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation.operandLayout(inputs.size, params) + @Suppress("UNCHECKED_CAST") + val lengths = layout.lengthsIndex?.let { inputs[it] as Tensor } listOf(ops.scaledDotProductAttention( query = inputs[0], key = inputs[1], value = inputs[2], - mask = inputs.getOrNull(3), + mask = layout.maskIndex?.let { inputs[it] }, scale = scale ?: 0f, - causal = causal + causal = causal, + keyLengths = lengths, + window = params["window"] as? Int ?: 0, + rightContext = params["rightContext"] as? Int ?: 0, )) } diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt index d938a9940..f1d07abd7 100644 --- a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt @@ -268,7 +268,7 @@ private class TestTensorOps : TensorOps { override fun sin(tensor: Tensor): Tensor = tensor override fun cos(tensor: Tensor): Tensor = tensor override fun convTranspose1d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Int, padding: Int, outputPadding: Int, dilation: Int, groups: Int): Tensor = input - override fun scaledDotProductAttention(query: Tensor, key: Tensor, value: Tensor, mask: Tensor?, scale: Float, causal: Boolean): Tensor = query + override fun scaledDotProductAttention(query: Tensor, key: Tensor, value: Tensor, mask: Tensor?, scale: Float, causal: Boolean, keyLengths: Tensor?, window: Int, rightContext: Int): Tensor = query override fun conv1d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Int, padding: Int, dilation: Int, groups: Int): Tensor = input override fun conv2d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Pair, padding: Pair, dilation: Pair, groups: Int): Tensor = input override fun conv3d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Triple, padding: Triple, dilation: Triple, groups: Int): Tensor = input diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/SdpaKeyLengthsRecordingTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/SdpaKeyLengthsRecordingTest.kt new file mode 100644 index 000000000..b536db8d4 --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/SdpaKeyLengthsRecordingTest.kt @@ -0,0 +1,110 @@ +package sk.ainet.exec + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import sk.ainet.exec.tensor.ops.DefaultCpuOps +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.graph.exec.ComputeGraphExecutor +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 + +/** + * #1343 acceptance: the mask rules of an eager (DIRECT) SDPA call are recorded as node metadata + * (`maskKind`, `window`, `rightContext`) with `keyLengths` as the last operand, and replaying the + * graph (OPTIMIZED) reproduces the eager output. + */ +class SdpaKeyLengthsRecordingTest { + + private val dataFactory = DenseTensorDataFactory() + + private fun tensor(shape: Shape, seed: Int): VoidOpsTensor { + val values = FloatArray(shape.volume) { i -> (((i * 37 + seed * 11) % 29) - 14) / 7f } + return VoidOpsTensor(dataFactory.fromFloatArray(shape, FP32::class, values), FP32::class) + } + + @Suppress("UNCHECKED_CAST") + private fun lengths(vararg values: Int): Tensor = + VoidOpsTensor(dataFactory.fromIntArray(Shape(values.size), Int32::class, values), Int32::class) as Tensor + + private fun floats(t: Tensor<*, *>): FloatArray = t.data.copyToFloatArray() + + @Test + fun key_lengths_are_recorded_as_metadata_and_replay_bit_identical() { + val ctx = DefaultGraphExecutionContext.tape(DefaultCpuOps(dataFactory)) + val q = tensor(Shape(2, 4, 6, 8), 1) + val k = tensor(Shape(2, 2, 6, 8), 2) + val v = tensor(Shape(2, 2, 6, 8), 3) + val w = tensor(Shape(8, 5), 4) + val len = lengths(6, 3) + + ctx.startRecording() + val ops = ctx.ops + val attn = ops.scaledDotProductAttention(q, k, v, keyLengths = len, window = 4, rightContext = 1) + val eager = ops.matmul(attn, w) + val tape = ctx.stopRecording() as DefaultExecutionTape + + // The trace carries the rules … + val trace = tape.traces.single { it.opType == "scaledDotProductAttention" } + assertEquals("combined", trace.attributes["maskKind"]) + assertEquals(4, trace.attributes["window"]) + assertEquals(1, trace.attributes["rightContext"]) + assertEquals(false, trace.attributes["causal"]) + assertEquals(4, trace.inputs.size, "q, k, v, keyLengths") + + // … and so does the graph node, with keyLengths as its last operand. + val graph = tape.toComputeGraph(synthesizeExternalInputs = true, embedConstants = false) + val node = graph.nodes.single { it.operation.name == "scaledDotProductAttention" } + assertEquals("combined", node.operation.parameters["maskKind"]) + val layout = ScaledDotProductAttentionOperation.operandLayout(node.inputs.size, node.operation.parameters) + assertEquals(ScaledDotProductAttentionOperation.OperandLayout(maskIndex = null, lengthsIndex = 3), layout) + + // Replay through the graph executor on the same kernels: external inputs keyed by trace id. + val session = ctx.session + @Suppress("UNCHECKED_CAST") + val inputs: Map> = mapOf( + session.refOf(q).id to q, session.refOf(k).id to k, session.refOf(v).id to v, + session.refOf(w).id to w, session.refOf(len).id to len, + ) as Map> + val outputs = ComputeGraphExecutor(graph, DefaultCpuOps(dataFactory)).execute(inputs) + val matmul = graph.nodes.single { it.operation.name == "matmul" } + val replayed = floats(outputs.getValue(matmul.id)) + val expected = floats(eager) + assertEquals(expected.size, replayed.size) + var maxDiff = 0f + for (i in expected.indices) maxDiff = maxOf(maxDiff, kotlin.math.abs(expected[i] - replayed[i])) + assertTrue(maxDiff <= 1e-6f, "DIRECT vs OPTIMIZED max |Δ| = $maxDiff") + + // Sanity: the rules really applied — the unmasked call differs on this block. + val plain = floats(DefaultCpuOps(dataFactory).scaledDotProductAttention(q, k, v)) + var ruleDiff = 0f + for (i in plain.indices) ruleDiff = maxOf(ruleDiff, kotlin.math.abs(plain[i] - floats(attn)[i])) + assertTrue(ruleDiff > 1e-5f, "masked and unmasked attention should differ, max |Δ| = $ruleDiff") + } + + @Test + fun dense_mask_and_lengths_keep_their_operand_slots() { + val ctx = DefaultGraphExecutionContext.tape(DefaultCpuOps(dataFactory)) + val q = tensor(Shape(1, 2, 3, 4), 1); val k = tensor(Shape(1, 2, 3, 4), 2); val v = tensor(Shape(1, 2, 3, 4), 3) + val mask = tensor(Shape(1, 1, 3, 3), 9) + ctx.startRecording() + ctx.ops.scaledDotProductAttention(q, k, v, mask = mask, causal = true, keyLengths = lengths(2)) + ctx.ops.scaledDotProductAttention(q, k, v, mask = mask, causal = true) + ctx.ops.scaledDotProductAttention(q, k, v, causal = true) + val tape = ctx.stopRecording() as DefaultExecutionTape + val traces = tape.traces.filter { it.opType == "scaledDotProductAttention" } + assertEquals(listOf("combined", "causal", "causal"), traces.map { it.attributes["maskKind"] }) + assertEquals(listOf(5, 4, 3), traces.map { it.inputs.size }) + val layouts = traces.map { ScaledDotProductAttentionOperation.operandLayout(it.inputs.size, it.attributes) } + assertEquals(ScaledDotProductAttentionOperation.OperandLayout(3, 4), layouts[0]) + assertEquals(ScaledDotProductAttentionOperation.OperandLayout(3, null), layouts[1]) + assertEquals(ScaledDotProductAttentionOperation.OperandLayout(null, null), layouts[2]) + } +} diff --git a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api index 81d59d78e..8f228dfd3 100644 --- a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api +++ b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api @@ -470,6 +470,14 @@ public final class sk/ainet/compile/hlo/converters/ScalarOperationsConverter : s public fun getSupportedOperations ()Ljava/util/Set; } +public final class sk/ainet/compile/hlo/converters/SdpaLowering { + public static final field INSTANCE Lsk/ainet/compile/hlo/converters/SdpaLowering; + public final fun effectiveScale (FI)F + public final fun emit (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/util/List;Ljava/util/List;Ljava/lang/String;Ljava/util/Map;)Ljava/lang/String; + public final fun parseTensorType (Ljava/lang/String;)Lkotlin/Pair; + public final fun renderType (Ljava/util/List;Ljava/lang/String;)Ljava/lang/String; +} + public final class sk/ainet/compile/hlo/converters/ShapeOperationsConverter : sk/ainet/compile/hlo/StableHloOperationConverter { public fun ()V public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult; diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt index 98163a63f..08b789c83 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt @@ -6,15 +6,14 @@ import sk.ainet.compile.hlo.StableHloOperationConverter import sk.ainet.lang.tensor.Dim import sk.ainet.lang.tensor.hasDynamic import sk.ainet.lang.graph.GraphNode -import kotlin.math.sqrt /** - * Converts `scaledDotProductAttention` to the standard StableHLO attention - * subgraph: + * Converts `scaledDotProductAttention` to the standard StableHLO attention subgraph through + * [SdpaLowering], the one builder of the formula (#1343): * * scores = Q · Kᵀ (dot_general, contract head_dim) * scaled = scores * scale (scale = arg, or 1/sqrt(head_dim)) - * attn = softmax(scaled, axis = -1) (max/sub/exp/sum/div, numerically stable) + * attn = softmax(scaled + mask) (max/sub/exp/sum/div, numerically stable) * out = attn · V (dot_general, contract key length) * * Q/K/V are batched `[.., S, D]`; batching dims are every leading dim except the @@ -23,10 +22,10 @@ import kotlin.math.sqrt * SDPA is a core `TensorOps` op (KSP-generated), so its converter lives here in * core alongside dot_general/softmax — the transformer modules just decompose to it. * - * Causal masking: when the `causal` attribute is set, an additive -inf mask - * (built from iota row/col indices + compare + select) is added to the scaled - * scores before softmax so each query only attends to keys at or before it. - * An explicit `mask` operand (operands[3]) takes priority over the iota path. + * Masking: the node's positional rules (`causal`, `window`/`rightContext`) and the `keyLengths` + * operand become an `iota`/`compare`/`select` keep mask, and a dense `mask` operand is added on + * top — the same rules the CPU kernel applies, defined once in `AttentionMask`. The rules need a + * static scores shape; a dynamic key length (`?`) is exported with a dense mask only. * * Grouped-query attention (SKEEP-005 phase 2, "structure at compile time"): when K/V * carry fewer heads than Q (`[b, nKV, Sk, hd]` vs `[b, H, Sq, hd]`, `nKV | H`), Q is @@ -65,16 +64,21 @@ public class AttentionOperationsConverter : StableHloOperationConverter { val outSpec = node.outputs.firstOrNull() val mapper = context.getTypeMapper() val elem = outSpec?.let { mapper.mapDType(it.dtype) } ?: "f32" - // Render a shape's dims, mapping a dynamic extent (DYNAMIC_DIM = -1) to `?`. - fun dims(shape: List): String = shape.joinToString("x") { Dim.render(it) } - fun typeOf(shape: List): String = "tensor<${dims(shape)}x$elem>" - - val qType = context.getValueType(operands[0]) ?: typeOf(qShape) - val kType = context.getValueType(operands[1]) ?: typeOf(kShape) - val vType = context.getValueType(operands[2]) ?: typeOf(vShape) - + fun typeOf(shape: List): String = SdpaLowering.renderType(shape, elem) + + // Operand MLIR types: the registered value type, else the node's input spec. The builder reads + // every shape it needs from these, as it must inside a composite decomposition function. + val operandTypes = operands.mapIndexed { i, operand -> + context.getValueType(operand) + ?: node.inputs.getOrNull(i)?.let { spec -> + val shape = spec.shape + if (i < 3 && shape != null) typeOf(shape) else mapper.mapTensorType(spec) + } + ?: when (i) { 0 -> typeOf(qShape); 1 -> typeOf(kShape); 2 -> typeOf(vShape) + else -> return ConversionResult.Failure("SDPA operand $i has no known type", "Missing operand type for ${node.id}") } + } val headDim = qShape[rank - 1] - val keyLen = kShape[rank - 2] + val outputType = outSpec?.let { mapper.mapTensorType(it) } ?: typeOf(qShape.dropLast(1) + headDim) // Grouped-query attention: K/V heads divide Q heads → work on Q as [b, nKV, nRep, Sq, hd]. val gqa = rank == 4 && kShape.size == 4 && kShape[1] != qShape[1] @@ -93,215 +97,40 @@ public class AttentionOperationsConverter : StableHloOperationConverter { ) } } - val qWork: List = if (gqa) listOf(qShape[0], kShape[1], qShape[1] / kShape[1], qShape[2], headDim) else qShape - val qWorkType = if (gqa) typeOf(qWork) else qType - val scoresShape = qWork.dropLast(1) + keyLen // [.., Sq, Sk] (GQA: [b, nKV, nRep, Sq, Sk]) - val scoresType = typeOf(scoresShape) - val outputType = outSpec?.let { mapper.mapTensorType(it) } ?: typeOf(qShape.dropLast(1) + headDim) - val outWorkType = if (gqa) typeOf(qWork.dropLast(1) + headDim) else outputType - - val scaleParam = (node.operation.parameters["scale"] as? Number)?.toFloat() ?: 0f - val scaleVal = if (scaleParam != 0f) scaleParam else (1.0f / sqrt(headDim.toFloat())) - - val hasBatch = rank > 2 - val batchList = (0 until rank - 2).joinToString(", ") // K/V batching dims; == Q's ([b, nKV]) under GQA too - val batchClause = if (hasBatch) "batching_dims = [$batchList] x [$batchList], " else "" - val contractQ = qWork.size - 1 // contract head_dim of Q … - val contractK = rank - 1 // … and K - val sdAxis = scoresShape.size - 1 // softmax over key length - val reducedShape = scoresShape.dropLast(1) - val reducedType = if (reducedShape.isEmpty()) "tensor<$elem>" else typeOf(reducedShape) - val bcastDims = (scoresShape.indices).filter { it != sdAxis }.joinToString(", ") - val contractAttn = scoresShape.size - 1 // attn key-length axis - val contractV = rank - 2 // V key-length axis - val causal = (node.operation.parameters["causal"] as? Boolean) ?: false - val qAxis = scoresShape.size - 2 // query position in scores [.., Sq, Sk] - val scoresI32Type = "tensor<${dims(scoresShape)}xi32>" - val scoresI1Type = "tensor<${dims(scoresShape)}xi1>" - - val scaleC = context.nextTempValue() - val scaleB = context.nextTempValue() - val qScaled = context.nextTempValue() - val scores = context.nextTempValue() - val maxInit = context.nextTempValue(); val maxV = context.nextTempValue(); val maxB = context.nextTempValue() - val shifted = context.nextTempValue(); val expV = context.nextTempValue() - val sumInit = context.nextTempValue(); val sumV = context.nextTempValue(); val sumB = context.nextTempValue() - val attn = context.nextTempValue() - val out = context.nextTempValue() - - val ops = mutableListOf( - // Scale is applied to Q *before* the QK dot — scores·s == (q·s)@kᵀ exactly. This keeps the scale - // constant at the STATIC query type (`[.., Sq, headDim]`) rather than a splat sized to the scores - // shape `[.., Sq, Sk]`, whose Sk (key/cache length) may be dynamic (`?`) in KV-cache decode — a - // dynamic-shape splat constant is invalid StableHLO. Also drops a full scores-sized dense constant - // from static graphs. - "$scaleC = stablehlo.constant dense<$scaleVal> : tensor<$elem>", - "$scaleB = stablehlo.broadcast_in_dim $scaleC, dims = [] : (tensor<$elem>) -> $qType", - "$qScaled = stablehlo.multiply ${operands[0]}, $scaleB : $qType", - ) - // GQA: expose the head groups as a batching dim of Q — a static, copy-free reshape. - val qForDot = if (gqa) { - val qg = context.nextTempValue() - ops += "$qg = stablehlo.reshape $qScaled : ($qType) -> $qWorkType" - qg - } else qScaled - ops += "$scores = stablehlo.dot_general $qForDot, ${operands[1]}, ${batchClause}contracting_dims = [$contractQ] x [$contractK] : ($qWorkType, $kType) -> $scoresType" - - // When the scores shape is dynamic (the `?` key/cache dim of KV-cache decode), everything broadcast onto - // it — the explicit mask and the softmax's reduced max/sum — must use `stablehlo.dynamic_broadcast_in_dim`: - // a static `stablehlo.broadcast_in_dim` cannot target a dynamic shape. Its runtime `output_dimensions` - // operand is built once, here, from the scores tensor via `get_dimension_size`. - // Static graphs keep the original explicit `broadcast_in_dim` path (byte-for-byte unchanged). - val dyn = scoresShape.hasDynamic() - val shapeType = "tensor<${scoresShape.size}xi32>" - val scoresShapeOperand: String = if (!dyn) "" else run { - val parts = scoresShape.indices.map { d -> - if (Dim.isStatic(scoresShape[d])) { - val c = context.nextTempValue() - ops += "$c = stablehlo.constant dense<${scoresShape[d]}> : tensor<1xi32>" - c - } else { - val gd = context.nextTempValue(); val gr = context.nextTempValue() - ops += "$gd = stablehlo.get_dimension_size $scores, dim = $d : ($scoresType) -> tensor" - ops += "$gr = stablehlo.reshape $gd : (tensor) -> tensor<1xi32>" - gr - } - } - val sh = context.nextTempValue() - ops += "$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> $shapeType" - sh - } - - // Explicit additive mask (operands[3]) — e.g. a sliding-window+causal - // mask the caller built and passed with causal=false. It already - // encodes causality/window, so it takes priority over the built-in - // iota causal path. Brought to the scores shape and added. Without this - // the masked layers run UNMASKED (attend to future tokens) — correct - // only at position 0. - // - // Shape rules: trailing-aligned broadcast. Under GQA a rank-4 mask [b, M, Sq, Sk] keeps its batch on - // scores dim 0 and skips the nRep axis (dim 2); M may be 1 (head-shared) or nKV (one row per group). - // A per-head mask (M = H) is first VIEWED as [b, nKV, nRep, Sq, Sk] in Q's own h = kv * nRep + r order - // (reshape, or dynamic_reshape when a dim is dynamic) — broadcasting H onto nKV is invalid IR. - var softmaxIn = scores // scores are already scaled (scale folded into Q above) - val maskOperand = operands.getOrNull(3) - if (maskOperand != null) { - var maskShape: List = node.inputs.getOrNull(3)?.shape ?: scoresShape - var maskVal = maskOperand - var maskType = context.getValueType(maskOperand) ?: typeOf(maskShape) - var maskDims: List = if (gqa && maskShape.size == 4) listOf(0, 1, 3, 4) else { - val offset = scoresShape.size - maskShape.size - maskShape.indices.map { it + offset } - } - if (gqa && maskShape.size == 4 && maskShape[1] != 1 && maskShape[1] != kShape[1]) { - val nKV = kShape[1]; val nH = qShape[1] - if (maskShape[1] != nH) { - return ConversionResult.Failure( - "SDPA grouped-query attention mask head dim must be 1, K/V heads ($nKV) or Q heads ($nH), got $maskShape", - "Unsupported GQA mask shape for ${node.id}", - ) - } - val split = listOf(maskShape[0], nKV, nH / nKV, maskShape[2], maskShape[3]) - val splitType = typeOf(split) - val r = context.nextTempValue() - if (!maskShape.hasDynamic()) { - ops += "$r = stablehlo.reshape $maskVal : ($maskType) -> $splitType" - } else { - // split dim -> source mask dim (the two group dims are static by construction) - val source = listOf(0, -1, -1, 2, 3) - val parts = split.indices.map { d -> - if (Dim.isStatic(split[d])) { - val c = context.nextTempValue() - ops += "$c = stablehlo.constant dense<${split[d]}> : tensor<1xi32>" - c - } else { - val gd = context.nextTempValue(); val gr = context.nextTempValue() - ops += "$gd = stablehlo.get_dimension_size $maskVal, dim = ${source[d]} : ($maskType) -> tensor" - ops += "$gr = stablehlo.reshape $gd : (tensor) -> tensor<1xi32>" - gr - } - } - val sh = context.nextTempValue() - ops += "$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> tensor<5xi32>" - ops += "$r = stablehlo.dynamic_reshape $maskVal, $sh : ($maskType, tensor<5xi32>) -> $splitType" - } - maskVal = r; maskShape = split; maskType = splitType; maskDims = split.indices.toList() - } - val maskBc = if (maskShape == scoresShape) { - maskVal - } else if (!dyn) { - val mb = context.nextTempValue() - ops += "$mb = stablehlo.broadcast_in_dim $maskVal, dims = [${maskDims.joinToString(", ")}] : ($maskType) -> $scoresType" - mb - } else { - // Dynamic target: dynamic_broadcast_in_dim, stating which operand dims expand. Without the - // hints a dynamic operand dim mapped onto a dynamic result dim is ambiguous (1 -> N or N -> N), - // and backends such as IREE refuse to lower it. A dynamic mask dim must match the scores dim - // it maps to (the key length), so it is non-expanding. Generic op syntax: the hint attributes - // are not part of every StableHLO version's pretty form. - val expanding = mutableListOf(); val nonExpanding = mutableListOf() - maskShape.indices.forEach { i -> - val o = maskShape[i]; val t = scoresShape[maskDims[i]] - when { - !Dim.isStatic(o) -> nonExpanding += i - Dim.isStatic(t) && o == t -> nonExpanding += i - Dim.isStatic(t) && o == 1 -> expanding += i - !Dim.isStatic(t) && o != 1 -> nonExpanding += i - } - } - fun arr(xs: List) = if (xs.isEmpty()) "array" else "array" - val mb = context.nextTempValue() - ops += "$mb = \"stablehlo.dynamic_broadcast_in_dim\"($maskVal, $scoresShapeOperand) <{broadcast_dimensions = ${arr(maskDims)}, " + - "known_expanding_dimensions = ${arr(expanding)}, known_nonexpanding_dimensions = ${arr(nonExpanding)}}> : " + - "($maskType, $shapeType) -> $scoresType" - mb - } - val masked = context.nextTempValue() - ops += "$masked = stablehlo.add $scores, $maskBc : $scoresType" - softmaxIn = masked - } else if (causal) { - val iotaQ = context.nextTempValue(); val iotaK = context.nextTempValue() - val keep = context.nextTempValue(); val zeros = context.nextTempValue() - val ninf = context.nextTempValue(); val maskAdd = context.nextTempValue() - val masked = context.nextTempValue() - ops += "$iotaQ = stablehlo.iota dim = $qAxis : $scoresI32Type" - ops += "$iotaK = stablehlo.iota dim = $sdAxis : $scoresI32Type" - ops += "$keep = stablehlo.compare GE, $iotaQ, $iotaK : ($scoresI32Type, $scoresI32Type) -> $scoresI1Type" - ops += "$zeros = stablehlo.constant dense<0.0> : $scoresType" - // Masked-fill with a large finite negative (not -inf): matches - // MultiHeadAttention.buildSlidingCausalMask (-1e30) and avoids a -inf splat in the - // masked-fill select, which can trip downstream greedy constant-folding. - ops += "$ninf = stablehlo.constant dense<-1.000000e+30> : $scoresType" - ops += "$maskAdd = stablehlo.select $keep, $zeros, $ninf : $scoresI1Type, $scoresType" - ops += "$masked = stablehlo.add $scores, $maskAdd : $scoresType" - softmaxIn = masked + val params = node.operation.parameters + val layout = sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation.operandLayout(operands.size, params) + val causal = (params["causal"] as? Boolean) ?: false + val window = (params["window"] as? Int) ?: 0 + val rightContext = (params["rightContext"] as? Int) ?: 0 + val scoresShape = qShape.dropLast(1) + kShape[rank - 2] + if ((causal || window > 0 || layout.lengthsIndex != null) && scoresShape.hasDynamic()) { + return ConversionResult.Failure( + "SDPA positional/padding rules (causal/window/keyLengths) need a static scores shape, got $scoresShape — pass a dense mask for a dynamic key length", + "Structured mask on a dynamic shape for ${node.id}", + ) } - - fun broadcastBack(src: String, dst: String) { - if (dyn) { - ops += "$dst = stablehlo.dynamic_broadcast_in_dim $src, $scoresShapeOperand, dims = [$bcastDims] : ($reducedType, $shapeType) -> $scoresType" - } else { - ops += "$dst = stablehlo.broadcast_in_dim $src, dims = [$bcastDims] : ($reducedType) -> $scoresType" + val maskIndex = layout.maskIndex + if (maskIndex != null && gqa) { + val maskShape = node.inputs.getOrNull(maskIndex)?.shape + if (maskShape != null && maskShape.size == 4 && maskShape[1] != 1 && maskShape[1] != kShape[1] && maskShape[1] != qShape[1]) { + return ConversionResult.Failure( + "SDPA grouped-query attention mask head dim must be 1, K/V heads (${kShape[1]}) or Q heads (${qShape[1]}), got $maskShape", + "Unsupported GQA mask shape for ${node.id}", + ) } } - ops += "$maxInit = stablehlo.constant dense<${mapper.negInfBits(elem)}> : tensor<$elem>" - ops += "$maxV = stablehlo.reduce($softmaxIn init: $maxInit) applies stablehlo.maximum across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType" - broadcastBack(maxV, maxB) - ops += "$shifted = stablehlo.subtract $softmaxIn, $maxB : $scoresType" - ops += "$expV = stablehlo.exponential $shifted : $scoresType" - ops += "$sumInit = stablehlo.constant dense<0.0> : tensor<$elem>" - ops += "$sumV = stablehlo.reduce($expV init: $sumInit) applies stablehlo.add across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType" - broadcastBack(sumV, sumB) - ops += "$attn = stablehlo.divide $expV, $sumB : $scoresType" - if (gqa) { - val outG = context.nextTempValue() - ops += "$outG = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outWorkType" - ops += "$out = stablehlo.reshape $outG : ($outWorkType) -> $outputType" - } else { - ops += "$out = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outputType" + val scaleParam = (params["scale"] as? Number)?.toFloat() ?: 0f + val lowered = params + mapOf( + "scale" to SdpaLowering.effectiveScale(scaleParam, headDim), + "causal" to causal, "window" to window, "rightContext" to rightContext, + ) + return try { + val emitter = ContextEmitter(context) + val out = SdpaLowering.emit(emitter, operands, operandTypes, outputType, lowered) + ConversionResult.Success(outputValueName = out, emittedOperations = emitter.lines) + } catch (e: IllegalArgumentException) { + ConversionResult.Failure(e.message ?: "SDPA lowering failed", "SDPA lowering failed for ${node.id}") } - ops.forEach { context.emitOperation(it) } - return ConversionResult.Success(outputValueName = out, emittedOperations = ops) } } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt new file mode 100644 index 000000000..21537589c --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt @@ -0,0 +1,326 @@ +package sk.ainet.compile.hlo.converters + +import sk.ainet.compile.hlo.TypeMapper +import sk.ainet.lang.tensor.Dim +import sk.ainet.lang.tensor.hasDynamic +import sk.ainet.lang.tensor.ops.AttentionMask +import sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation + +/** + * The `scaledDotProductAttention` decomposition (#1343), written once through an [HloEmitter]: + * + * scores = (Q · scale) · Kᵀ (dot_general, contract head_dim) + * masked = scores + rules + mask (iota/compare/select from the recorded rules, dense mask added) + * attn = softmax(masked, axis = -1) (max/sub/exp/sum/div, numerically stable) + * out = attn · V (dot_general, contract key length) + * + * Everything the builder needs is in the operand MLIR types and the node parameters, so the same + * code produces the inline form and the body of a `stablehlo.composite` decomposition function. + * + * **Mask rules** ([AttentionMask]): `causal`, a sliding `window` with `rightContext`, and key + * padding from the `keyLengths` operand (`[batch]` integers) become one `i1` keep mask built with + * `stablehlo.iota` over the query and key axes of the scores, `compare`, `and`, then + * `select keep, 0, -1e30` added to the scores — no `[batch, 1, Sq, Sk]` constant enters the + * module. The causal diagonal is offset by `Sk − Sq` like the CPU kernel, so a KV-cache decode + * step attends to its whole cache. A dense `mask` operand is added on top, as the CPU kernel does. + * Positional rules need static scores; a dynamic key length (`?`) is supported with a dense mask only. + * + * **Grouped-query attention** (SKEEP-005 phase 2): when K/V carry fewer heads than Q, Q is viewed + * as `[b, nKV, nRep, Sq, hd]` and both `dot_general`s batch over `[b, nKV]`; K/V are never tiled. + */ +public object SdpaLowering { + + /** `tensor<1x2x?x8xf32>` → dims `[1, 2, -1, 8]` and element type `f32`; `tensor` → `[]`. */ + public fun parseTensorType(type: String): Pair, String> { + require(type.startsWith("tensor<") && type.endsWith(">")) { "not a tensor type: $type" } + val parts = type.removePrefix("tensor<").removeSuffix(">").split('x') + val elem = parts.last() + val dims = parts.dropLast(1).map { if (it == "?") TypeMapper.DYNAMIC_DIM else it.toInt() } + return dims to elem + } + + public fun renderType(dims: List, elem: String): String = + if (dims.isEmpty()) "tensor<$elem>" else "tensor<${dims.joinToString("x") { Dim.render(it) }}x$elem>" + + /** The scale the formula uses: the recorded value, or `1/√headDim` when it was left at `0`. */ + public fun effectiveScale(scale: Float, headDim: Int): Float = + if (scale != 0f) scale else (1.0f / kotlin.math.sqrt(headDim.toFloat())) + + /** + * Emit the decomposition for [operands] (`q, k, v[, mask][, keyLengths]`, see + * [ScaledDotProductAttentionOperation.operandLayout]) of MLIR [operandTypes] producing + * [resultType]; returns the result value. [params] are the node parameters with `scale` + * already resolved ([effectiveScale]). Shapes are validated by the converter beforehand. + */ + public fun emit( + e: HloEmitter, + operands: List, + operandTypes: List, + resultType: String, + params: Map, + ): String { + require(operands.size >= 3 && operands.size == operandTypes.size) { "SDPA needs q, k, v operands with types" } + val layout = ScaledDotProductAttentionOperation.operandLayout(operands.size, params) + val (qShape, elem) = parseTensorType(operandTypes[0]) + val (kShape, _) = parseTensorType(operandTypes[1]) + val qType = operandTypes[0]; val kType = operandTypes[1]; val vType = operandTypes[2] + val rank = qShape.size + val headDim = qShape[rank - 1] + val keyLen = kShape[rank - 2] + val outputType = resultType + fun typeOf(shape: List) = renderType(shape, elem) + + val gqa = rank == 4 && kShape.size == 4 && kShape[1] != qShape[1] + val qWork: List = if (gqa) listOf(qShape[0], kShape[1], qShape[1] / kShape[1], qShape[2], headDim) else qShape + val qWorkType = if (gqa) typeOf(qWork) else qType + val scoresShape = qWork.dropLast(1) + keyLen // [.., Sq, Sk] (GQA: [b, nKV, nRep, Sq, Sk]) + val scoresType = typeOf(scoresShape) + val outWorkType = if (gqa) typeOf(qWork.dropLast(1) + headDim) else outputType + + val scaleVal = effectiveScale((params["scale"] as? Number)?.toFloat() ?: 0f, headDim) + val causal = (params["causal"] as? Boolean) ?: false + val window = (params["window"] as? Int) ?: 0 + val rightContext = (params["rightContext"] as? Int) ?: 0 + AttentionMask.validate(window, rightContext) + + val hasBatch = rank > 2 + val batchList = (0 until rank - 2).joinToString(", ") // K/V batching dims; == Q's ([b, nKV]) under GQA too + val batchClause = if (hasBatch) "batching_dims = [$batchList] x [$batchList], " else "" + val contractQ = qWork.size - 1 // contract head_dim of Q … + val contractK = rank - 1 // … and K + val sdAxis = scoresShape.size - 1 // softmax over key length + val qAxis = scoresShape.size - 2 // query position in scores [.., Sq, Sk] + val reducedShape = scoresShape.dropLast(1) + val reducedType = if (reducedShape.isEmpty()) "tensor<$elem>" else typeOf(reducedShape) + val bcastDims = (scoresShape.indices).filter { it != sdAxis }.joinToString(", ") + val contractAttn = scoresShape.size - 1 // attn key-length axis + val contractV = rank - 2 // V key-length axis + val scoresI32Type = "tensor<${scoresShape.joinToString("x") { Dim.render(it) }}xi32>" + val scoresI1Type = "tensor<${scoresShape.joinToString("x") { Dim.render(it) }}xi1>" + + // Scale is applied to Q *before* the QK dot — scores·s == (q·s)@kᵀ exactly. This keeps the scale + // constant at the STATIC query type (`[.., Sq, headDim]`) rather than a splat sized to the scores + // shape `[.., Sq, Sk]`, whose Sk (key/cache length) may be dynamic (`?`) in KV-cache decode — a + // dynamic-shape splat constant is invalid StableHLO. Also drops a full scores-sized dense constant + // from static graphs. + val scaleC = e.newValue(); val scaleB = e.newValue(); val qScaled = e.newValue(); val scores = e.newValue() + e.op("$scaleC = stablehlo.constant dense<$scaleVal> : tensor<$elem>") + e.op("$scaleB = stablehlo.broadcast_in_dim $scaleC, dims = [] : (tensor<$elem>) -> $qType") + e.op("$qScaled = stablehlo.multiply ${operands[0]}, $scaleB : $qType") + // GQA: expose the head groups as a batching dim of Q — a static, copy-free reshape. + val qForDot = if (gqa) { + val qg = e.newValue() + e.op("$qg = stablehlo.reshape $qScaled : ($qType) -> $qWorkType") + qg + } else qScaled + e.op("$scores = stablehlo.dot_general $qForDot, ${operands[1]}, ${batchClause}contracting_dims = [$contractQ] x [$contractK] : ($qWorkType, $kType) -> $scoresType") + + // When the scores shape is dynamic (the `?` key/cache dim of KV-cache decode), everything broadcast onto + // it — the explicit mask and the softmax's reduced max/sum — must use `stablehlo.dynamic_broadcast_in_dim`: + // a static `stablehlo.broadcast_in_dim` cannot target a dynamic shape. Its runtime `output_dimensions` + // operand is built once, here, from the scores tensor via `get_dimension_size`. + // Static graphs keep the original explicit `broadcast_in_dim` path (byte-for-byte unchanged). + val dyn = scoresShape.hasDynamic() + val shapeType = "tensor<${scoresShape.size}xi32>" + val scoresShapeOperand: String = if (!dyn) "" else run { + val parts = scoresShape.indices.map { d -> + if (Dim.isStatic(scoresShape[d])) { + val c = e.newValue() + e.op("$c = stablehlo.constant dense<${scoresShape[d]}> : tensor<1xi32>") + c + } else { + val gd = e.newValue(); val gr = e.newValue() + e.op("$gd = stablehlo.get_dimension_size $scores, dim = $d : ($scoresType) -> tensor") + e.op("$gr = stablehlo.reshape $gd : (tensor) -> tensor<1xi32>") + gr + } + } + val sh = e.newValue() + e.op("$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> $shapeType") + sh + } + + var softmaxIn = scores // scores are already scaled (scale folded into Q above) + + // Explicit additive mask — e.g. a sliding-window+causal mask the caller built, or a cross-attention + // padding mask. Brought to the scores shape and added. Without this the masked layers run UNMASKED + // (attend to future tokens) — correct only at position 0. + // + // Shape rules: trailing-aligned broadcast. Under GQA a rank-4 mask [b, M, Sq, Sk] keeps its batch on + // scores dim 0 and skips the nRep axis (dim 2); M may be 1 (head-shared) or nKV (one row per group). + // A per-head mask (M = H) is first VIEWED as [b, nKV, nRep, Sq, Sk] in Q's own h = kv * nRep + r order + // (reshape, or dynamic_reshape when a dim is dynamic) — broadcasting H onto nKV is invalid IR. + val maskOperand = layout.maskIndex?.let { operands[it] } + if (maskOperand != null) { + var maskType = operandTypes[layout.maskIndex!!] + var maskShape: List = parseTensorType(maskType).first + var maskVal = maskOperand + var maskDims: List = if (gqa && maskShape.size == 4) listOf(0, 1, 3, 4) else { + val offset = scoresShape.size - maskShape.size + maskShape.indices.map { it + offset } + } + if (gqa && maskShape.size == 4 && maskShape[1] != 1 && maskShape[1] != kShape[1]) { + val nKV = kShape[1]; val nH = qShape[1] + require(maskShape[1] == nH) { + "SDPA grouped-query attention mask head dim must be 1, K/V heads ($nKV) or Q heads ($nH), got $maskShape" + } + val split = listOf(maskShape[0], nKV, nH / nKV, maskShape[2], maskShape[3]) + val splitType = typeOf(split) + val r = e.newValue() + if (!maskShape.hasDynamic()) { + e.op("$r = stablehlo.reshape $maskVal : ($maskType) -> $splitType") + } else { + // split dim -> source mask dim (the two group dims are static by construction) + val source = listOf(0, -1, -1, 2, 3) + val parts = split.indices.map { d -> + if (Dim.isStatic(split[d])) { + val c = e.newValue() + e.op("$c = stablehlo.constant dense<${split[d]}> : tensor<1xi32>") + c + } else { + val gd = e.newValue(); val gr = e.newValue() + e.op("$gd = stablehlo.get_dimension_size $maskVal, dim = ${source[d]} : ($maskType) -> tensor") + e.op("$gr = stablehlo.reshape $gd : (tensor) -> tensor<1xi32>") + gr + } + } + val sh = e.newValue() + e.op("$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> tensor<5xi32>") + e.op("$r = stablehlo.dynamic_reshape $maskVal, $sh : ($maskType, tensor<5xi32>) -> $splitType") + } + maskVal = r; maskShape = split; maskType = splitType; maskDims = split.indices.toList() + } + val maskBc = if (maskShape == scoresShape) { + maskVal + } else if (!dyn) { + val mb = e.newValue() + e.op("$mb = stablehlo.broadcast_in_dim $maskVal, dims = [${maskDims.joinToString(", ")}] : ($maskType) -> $scoresType") + mb + } else { + // Dynamic target: dynamic_broadcast_in_dim, stating which operand dims expand. Without the + // hints a dynamic operand dim mapped onto a dynamic result dim is ambiguous (1 -> N or N -> N), + // and backends such as IREE refuse to lower it. A dynamic mask dim must match the scores dim + // it maps to (the key length), so it is non-expanding. Generic op syntax: the hint attributes + // are not part of every StableHLO version's pretty form. + val expanding = mutableListOf(); val nonExpanding = mutableListOf() + maskShape.indices.forEach { i -> + val o = maskShape[i]; val t = scoresShape[maskDims[i]] + when { + !Dim.isStatic(o) -> nonExpanding += i + Dim.isStatic(t) && o == t -> nonExpanding += i + Dim.isStatic(t) && o == 1 -> expanding += i + !Dim.isStatic(t) && o != 1 -> nonExpanding += i + } + } + fun arr(xs: List) = if (xs.isEmpty()) "array" else "array" + val mb = e.newValue() + e.op("$mb = \"stablehlo.dynamic_broadcast_in_dim\"($maskVal, $scoresShapeOperand) <{broadcast_dimensions = ${arr(maskDims)}, " + + "known_expanding_dimensions = ${arr(expanding)}, known_nonexpanding_dimensions = ${arr(nonExpanding)}}> : " + + "($maskType, $shapeType) -> $scoresType") + mb + } + val masked = e.newValue() + e.op("$masked = stablehlo.add $softmaxIn, $maskBc : $scoresType") + softmaxIn = masked + } + + // Positional and padding rules (#1343) as one i1 keep mask over the scores shape. Each term is the + // AttentionMask rule written with iota/compare; terms are and-ed, then masked-filled with a large + // finite negative (not -inf): matches AttentionMask.MASKED / the dense masks callers build and + // avoids a -inf splat in the select, which can trip downstream greedy constant-folding. + val lengthsOperand = layout.lengthsIndex?.let { operands[it] } + if (causal || window > 0 || lengthsOperand != null) { + require(!dyn) { "SDPA positional/padding rules need a static scores shape, got $scoresShape — pass a dense mask for a dynamic key length" } + val seqQ = scoresShape[qAxis]; val seqKV = scoresShape[sdAxis] + val iotaQ = e.newValue(); val iotaK = e.newValue() + e.op("$iotaQ = stablehlo.iota dim = $qAxis : $scoresI32Type") + e.op("$iotaK = stablehlo.iota dim = $sdAxis : $scoresI32Type") + // absQ = Sk - Sq + qi: the last query aligns with the last key (AttentionMask.queryPosition). + val offset = seqKV - seqQ + val qPos = if (offset == 0) iotaQ else shiftedI32(e, iotaQ, offset, scoresI32Type) + val terms = mutableListOf() + if (causal) { + val keep = e.newValue() + e.op("$keep = stablehlo.compare GE, $qPos, $iotaK : ($scoresI32Type, $scoresI32Type) -> $scoresI1Type") + terms += keep + } + if (window > 0) { + val lower = shiftedI32(e, qPos, 1 - window, scoresI32Type) + val keepLower = e.newValue() + e.op("$keepLower = stablehlo.compare GE, $iotaK, $lower : ($scoresI32Type, $scoresI32Type) -> $scoresI1Type") + terms += keepLower + if (!causal) { + // The band's right edge; under `causal` the diagonal already bounds it (rightContext <= 0 there). + val upper = if (rightContext == 0) qPos else shiftedI32(e, qPos, rightContext, scoresI32Type) + val keepUpper = e.newValue() + e.op("$keepUpper = stablehlo.compare LE, $iotaK, $upper : ($scoresI32Type, $scoresI32Type) -> $scoresI1Type") + terms += keepUpper + } + } + if (lengthsOperand != null) { + val lengthsType = operandTypes[layout.lengthsIndex!!] + val (lenShape, lenElem) = parseTensorType(lengthsType) + require(lenShape.size == 1 && lenShape[0] == scoresShape[0]) { "SDPA keyLengths must be [batch=${scoresShape[0]}], got $lengthsType" } + var lengths = lengthsOperand + if (lenElem != "i32") { + val conv = e.newValue() + e.op("$conv = stablehlo.convert $lengths : ($lengthsType) -> tensor<${lenShape[0]}xi32>") + lengths = conv + } + val lenB = e.newValue(); val keepLen = e.newValue() + e.op("$lenB = stablehlo.broadcast_in_dim $lengths, dims = [0] : (tensor<${lenShape[0]}xi32>) -> $scoresI32Type") + e.op("$keepLen = stablehlo.compare LT, $iotaK, $lenB : ($scoresI32Type, $scoresI32Type) -> $scoresI1Type") + terms += keepLen + } + var keep = terms[0] + for (t in terms.drop(1)) { + val both = e.newValue() + e.op("$both = stablehlo.and $keep, $t : $scoresI1Type") + keep = both + } + val zeros = e.newValue(); val ninf = e.newValue(); val maskAdd = e.newValue(); val masked = e.newValue() + e.op("$zeros = stablehlo.constant dense<0.0> : $scoresType") + e.op("$ninf = stablehlo.constant dense<-1.000000e+30> : $scoresType") + e.op("$maskAdd = stablehlo.select $keep, $zeros, $ninf : $scoresI1Type, $scoresType") + e.op("$masked = stablehlo.add $softmaxIn, $maskAdd : $scoresType") + softmaxIn = masked + } + + fun broadcastBack(src: String, dst: String) { + if (dyn) { + e.op("$dst = stablehlo.dynamic_broadcast_in_dim $src, $scoresShapeOperand, dims = [$bcastDims] : ($reducedType, $shapeType) -> $scoresType") + } else { + e.op("$dst = stablehlo.broadcast_in_dim $src, dims = [$bcastDims] : ($reducedType) -> $scoresType") + } + } + val maxInit = e.newValue(); val maxV = e.newValue(); val maxB = e.newValue() + val shifted = e.newValue(); val expV = e.newValue() + val sumInit = e.newValue(); val sumV = e.newValue(); val sumB = e.newValue() + val attn = e.newValue(); val out = e.newValue() + e.op("$maxInit = stablehlo.constant dense<${TypeMapper().negInfBits(elem)}> : tensor<$elem>") + e.op("$maxV = stablehlo.reduce($softmaxIn init: $maxInit) applies stablehlo.maximum across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType") + broadcastBack(maxV, maxB) + e.op("$shifted = stablehlo.subtract $softmaxIn, $maxB : $scoresType") + e.op("$expV = stablehlo.exponential $shifted : $scoresType") + e.op("$sumInit = stablehlo.constant dense<0.0> : tensor<$elem>") + e.op("$sumV = stablehlo.reduce($expV init: $sumInit) applies stablehlo.add across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType") + broadcastBack(sumV, sumB) + e.op("$attn = stablehlo.divide $expV, $sumB : $scoresType") + if (gqa) { + val outG = e.newValue() + e.op("$outG = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outWorkType") + e.op("$out = stablehlo.reshape $outG : ($outWorkType) -> $outputType") + } else { + e.op("$out = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outputType") + } + return out + } + + /** `value + delta` as an i32 splat add over the scores shape. */ + private fun shiftedI32(e: HloEmitter, value: String, delta: Int, i32Type: String): String { + val c = e.newValue(); val r = e.newValue() + e.op("$c = stablehlo.constant dense<$delta> : $i32Type") + e.op("$r = stablehlo.add $value, $c : $i32Type") + return r + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/SdpaMaskRulesHloTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/SdpaMaskRulesHloTest.kt new file mode 100644 index 000000000..90a63bccd --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/SdpaMaskRulesHloTest.kt @@ -0,0 +1,102 @@ +package sk.ainet.compile.hlo + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue +import sk.ainet.compile.hlo.validation.MlirParser +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.DType + +/** + * #1343: the SDPA mask rules (`causal`, `window`/`rightContext`, `keyLengths`) lower to an + * `iota`/`compare`/`select` keep mask built from the node metadata — no `[b, 1, Sq, Sk]` + * constant or graph input — and the causal diagonal carries the `Sk − Sq` offset the CPU kernel uses. + */ +class SdpaMaskRulesHloTest { + + private fun graph( + q: List, kv: List, params: Map, mask: List? = null, lengths: Int? = null, + ): DefaultComputeGraph { + val g = DefaultComputeGraph() + fun input(id: String, shape: List, dtype: String = "FP32") = + GraphNode(id, InputOperation(), emptyList(), listOf(TensorSpec(id, shape, dtype))).also { g.addNode(it) } + val qs = input("q", q); val ks = input("k", kv); val vs = input("v", kv) + val ms = mask?.let { input("m", it) } + val ls = lengths?.let { input("len", listOf(it), "Int32") } + val operands = listOfNotNull(qs, ks, vs, ms, ls) + val out = TensorSpec("o", q, "FP32") + val sdpa = GraphNode("sdpa", ScaledDotProductAttentionOperation(params), operands.map { it.outputs[0] }, listOf(out)) + g.addNode(sdpa) + operands.forEachIndexed { i, n -> g.addEdge(GraphEdge("e$i", n, sdpa, 0, i, n.outputs[0])) } + return g + } + + private fun convert(g: DefaultComputeGraph, name: String = "t") = StableHloConverterFactory.createBasic().convert(g, name).content + + private fun params(causal: Boolean = false, window: Int = 0, rightContext: Int = 0, lengths: Boolean = false) = + ScaledDotProductAttentionOperation.parameters(0f, causal, window, rightContext, lengths) + + @Test + fun key_lengths_lower_to_iota_compare_select_from_the_lengths_operand() { + val mlir = convert(graph(listOf(3, 4, 8, 16), listOf(3, 4, 8, 16), params(lengths = true), lengths = 3)) + assertTrue(mlir.contains("%arg3: tensor<3xi32>"), "lengths are a plain [batch] i32 input:\n$mlir") + assertTrue(mlir.contains("stablehlo.broadcast_in_dim %arg3, dims = [0] : (tensor<3xi32>) -> tensor<3x4x8x8xi32>"), "lengths broadcast over the scores:\n$mlir") + assertTrue(mlir.contains("stablehlo.iota dim = 3 : tensor<3x4x8x8xi32>"), "key iota:\n$mlir") + assertTrue(mlir.contains("stablehlo.compare LT, %"), "ki < length:\n$mlir") + assertTrue(mlir.contains("stablehlo.select %") && mlir.contains("dense<-1.000000e+30> : tensor<3x4x8x8xf32>"), "masked fill:\n$mlir") + assertFalse(mlir.contains("%arg3: tensor<3x1x8x8xf32>"), "no dense mask input for padding:\n$mlir") + assertTrue(MlirValidator().validate(mlir).isEmpty(), MlirValidator().validate(mlir).toString()) + assertEquals("t", MlirParser().parse(mlir).getOrThrow().functionName) + } + + @Test + fun window_band_and_causal_offset_follow_the_shared_rules() { + // Chunked prefill: Sq = 2 over Sk = 6 → absQ = 4 + qi; window 3 → lower = absQ - 2. + val mlir = convert(graph(listOf(1, 2, 2, 8), listOf(1, 2, 6, 8), params(causal = true, window = 3))) + assertTrue(mlir.contains("stablehlo.constant dense<4> : tensor<1x2x2x6xi32>"), "causal offset Sk - Sq = 4:\n$mlir") + assertTrue(mlir.contains("stablehlo.compare GE, %"), "causal diagonal:\n$mlir") + assertTrue(mlir.contains("stablehlo.constant dense<-2> : tensor<1x2x2x6xi32>"), "window lower edge 1 - window:\n$mlir") + assertTrue(mlir.contains("stablehlo.and %"), "rules intersect:\n$mlir") + assertFalse(mlir.contains("compare LE"), "under causal the band has no separate right edge:\n$mlir") + + // Bidirectional band: window 4, rightContext 2, no causal → both edges, no diagonal. + val band = convert(graph(listOf(1, 2, 6, 8), listOf(1, 2, 6, 8), params(window = 4, rightContext = 2))) + assertTrue(band.contains("dense<-3> : tensor<1x2x6x6xi32>") && band.contains("dense<2> : tensor<1x2x6x6xi32>"), "band edges:\n$band") + assertTrue(band.contains("compare GE") && band.contains("compare LE"), "both edges:\n$band") + assertFalse(band.contains("dense<0> : tensor<1x2x6x6xi32>"), "no zero offset constant when Sq == Sk:\n$band") + } + + @Test + fun causal_prefill_keeps_the_previous_emission() { + // Sq == Sk: offset 0 is skipped, so the lines are what the converter emitted before #1343. + val mlir = convert(graph(listOf(1, 1, 2, 4), listOf(1, 1, 2, 4), mapOf("scale" to 0.0f, "causal" to true))) + assertTrue(mlir.contains("stablehlo.iota dim = 2 : tensor<1x1x2x2xi32>")) + assertTrue(mlir.contains("stablehlo.iota dim = 3 : tensor<1x1x2x2xi32>")) + assertTrue(Regex("stablehlo\\.compare GE, %v\\d+, %v\\d+ : \\(tensor<1x1x2x2xi32>, tensor<1x1x2x2xi32>\\) -> tensor<1x1x2x2xi1>").containsMatchIn(mlir), mlir) + assertFalse(mlir.contains("stablehlo.and"), "a single rule needs no and:\n$mlir") + assertFalse(mlir.contains("stablehlo.constant dense<0> : tensor<1x1x2x2xi32>"), "no offset add:\n$mlir") + } + + @Test + fun dense_mask_and_rules_combine_like_the_cpu_kernel() { + val mlir = convert(graph(listOf(2, 2, 4, 8), listOf(2, 2, 4, 8), params(causal = true, lengths = true), mask = listOf(2, 1, 4, 4), lengths = 2)) + assertTrue(mlir.contains("stablehlo.broadcast_in_dim %arg3, dims = [0, 1, 2, 3] : (tensor<2x1x4x4xf32>) -> tensor<2x2x4x4xf32>"), "dense mask added:\n$mlir") + assertTrue(mlir.contains("stablehlo.broadcast_in_dim %arg4, dims = [0] : (tensor<2xi32>) -> tensor<2x2x4x4xi32>"), "lengths after the mask:\n$mlir") + assertTrue(mlir.contains("compare GE") && mlir.contains("compare LT"), "causal and padding both apply:\n$mlir") + } + + @Test + fun rules_on_a_dynamic_key_length_are_refused_but_a_dense_mask_is_not() { + val dyn = listOf(1, 2, TypeMapper.DYNAMIC_DIM, 8) + val refused = runCatching { convert(graph(listOf(1, 2, 1, 8), dyn, params(lengths = true), lengths = 1)) } + assertTrue(refused.isFailure || !refused.getOrThrow().contains("stablehlo.iota"), "structured rules need static scores") + val dense = convert(graph(listOf(1, 2, 1, 8), dyn, params(), mask = listOf(1, 1, 1, TypeMapper.DYNAMIC_DIM))) + assertTrue(dense.contains("stablehlo.dynamic_broadcast_in_dim"), "dense mask on the dynamic path:\n$dense") + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 86868f9af..557860fc1 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -6247,6 +6247,25 @@ public final class sk/ainet/lang/tensor/ops/AddOperation : sk/ainet/lang/tensor/ public fun validateInputs (Ljava/util/List;)Lsk/ainet/lang/tensor/ops/ValidationResult; } +public final class sk/ainet/lang/tensor/ops/AttentionMask { + public static final field INSTANCE Lsk/ainet/lang/tensor/ops/AttentionMask; + public static final field KIND_CAUSAL Ljava/lang/String; + public static final field KIND_COMBINED Ljava/lang/String; + public static final field KIND_KEY_PADDING Ljava/lang/String; + public static final field KIND_NONE Ljava/lang/String; + public static final field KIND_SLIDING_WINDOW Ljava/lang/String; + public static final field MASKED F + public final fun additive (IIIZII[I)[F + public final fun attends (IIIIZIII)Z + public final fun endKey (IIIZIII)I + public final fun firstKey (IIII)I + public final fun hasLengths (Ljava/lang/String;)Z + public final fun keyRange (IIIZIII)Lkotlin/ranges/IntRange; + public final fun kindOf (ZIZ)Ljava/lang/String; + public final fun queryPosition (III)I + public final fun validate (II)V +} + public abstract class sk/ainet/lang/tensor/ops/BaseOperation : sk/ainet/lang/tensor/ops/Operation { public fun (Ljava/lang/String;Ljava/lang/String;Ljava/util/Map;)V public synthetic fun (Ljava/lang/String;Ljava/lang/String;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -6482,7 +6501,7 @@ public final class sk/ainet/lang/tensor/ops/KspTensorOps : sk/ainet/lang/tensor/ public fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)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; @@ -6637,6 +6656,7 @@ public final class sk/ainet/lang/tensor/ops/ReshapeOperation : sk/ainet/lang/ten } public final class sk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation : sk/ainet/lang/tensor/ops/BaseOperation { + public static final field Companion Lsk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation$Companion; public fun ()V public fun (Ljava/util/Map;)V public synthetic fun (Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -6646,6 +6666,24 @@ public final class sk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation : public fun validateInputs (Ljava/util/List;)Lsk/ainet/lang/tensor/ops/ValidationResult; } +public final class sk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation$Companion { + public final fun operandLayout (ILjava/util/Map;)Lsk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation$OperandLayout; + public final fun parameters (FZIIZ)Ljava/util/Map; +} + +public final class sk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation$OperandLayout { + public fun (Ljava/lang/Integer;Ljava/lang/Integer;)V + public final fun component1 ()Ljava/lang/Integer; + public final fun component2 ()Ljava/lang/Integer; + public final fun copy (Ljava/lang/Integer;Ljava/lang/Integer;)Lsk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation$OperandLayout; + public static synthetic fun copy$default (Lsk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation$OperandLayout;Ljava/lang/Integer;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/tensor/ops/ScaledDotProductAttentionOperation$OperandLayout; + public fun equals (Ljava/lang/Object;)Z + public final fun getLengthsIndex ()Ljava/lang/Integer; + public final fun getMaskIndex ()Ljava/lang/Integer; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/tensor/ops/SigmoidOperation : sk/ainet/lang/tensor/ops/BaseOperation { public fun ()V public fun (Ljava/util/Map;)V @@ -6757,8 +6795,8 @@ public abstract interface class sk/ainet/lang/tensor/ops/TensorOps { public abstract fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor; public abstract fun rsubScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; - public abstract 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 static synthetic fun scaledDotProductAttention$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;FZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public abstract 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 static synthetic fun scaledDotProductAttention$default (Lsk/ainet/lang/tensor/ops/TensorOps;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;IIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun sigmoid (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun sign (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun silu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; @@ -6802,7 +6840,7 @@ public final class sk/ainet/lang/tensor/ops/TensorOps$DefaultImpls { public static synthetic fun maxPool2d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun mean$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; - public static synthetic fun scaledDotProductAttention$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;FZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public static synthetic fun scaledDotProductAttention$default (Lsk/ainet/lang/tensor/ops/TensorOps;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;IIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun softmax$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun squeeze$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun sum$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; @@ -6962,7 +7000,7 @@ public final class sk/ainet/lang/tensor/ops/VoidTensorOps : sk/ainet/lang/tensor public fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)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; @@ -8066,6 +8104,7 @@ public final class sk/ainet/lang/trace/OpAttributeFactory { public final fun binary (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Ljava/util/Map; public final fun conv2d (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;I)Ljava/util/Map; public final fun scalarOp (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;Z)Ljava/util/Map; + public final fun sdpaMaskKind (Lsk/ainet/lang/tensor/Tensor;ZI)Ljava/util/Map; public final fun shapesAndDTypes (Ljava/util/List;Ljava/util/List;)Ljava/util/Map; public final fun unary (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Ljava/util/Map; } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/AttentionMask.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/AttentionMask.kt new file mode 100644 index 000000000..8ef20e59a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/AttentionMask.kt @@ -0,0 +1,104 @@ +package sk.ainet.lang.tensor.ops + +/** + * The one definition of which keys a query row of [TensorOps.scaledDotProductAttention] attends to (#1343). + * + * Every mask the op knows is a *structural* rule over positions, not a tensor: + * + * - **causal**: key `ki` is attended when `ki <= absQ`; + * - **sliding window** of size `w` with `rightContext = r`: `absQ - w + 1 <= ki <= absQ + r` + * (`r = 0` is the classic causal left window, `r > 0` a local bidirectional band); + * - **key padding** with per-row `length`: `ki < length`. + * + * `absQ = seqKV - seqQ + qi` aligns the last query with the last key, which is what a KV-cache + * decode step needs (`seqQ = 1`, `seqKV = pos + 1`) and reduces to `qi` for `seqQ == seqKV`. + * When several rules apply their key sets intersect. A masked score is `score + MASKED` + * ([MASKED] is a large finite negative, never `-inf`), so a fully masked row stays finite. + * + * The CPU kernel, the StableHLO decomposition and the dense fallback mask ([additive]) all read + * this object, so the three cannot drift; the kind of rule in force travels on the graph node as + * `parameters["maskKind"]` ([kindOf]) and `keyLengths` as its own operand, so a compiled target + * can regenerate the cheapest legal form (lengths for a fused kernel, iota/compare/select for a + * portable decomposition) instead of inheriting a `[batch, 1, seqQ, seqKV]` constant. + */ +public object AttentionMask { + + /** Additive value of a masked score: `-1e30`, the value the attention paths agreed on before this builder existed. */ + public const val MASKED: Float = -1.0e30f + + public const val KIND_NONE: String = "none" + public const val KIND_CAUSAL: String = "causal" + public const val KIND_SLIDING_WINDOW: String = "slidingWindow" + public const val KIND_KEY_PADDING: String = "keyPadding" + public const val KIND_COMBINED: String = "combined" + + /** The recorded `maskKind` for a call: padding alone, a positional rule alone, both (`combined`), or `none`. */ + public fun kindOf(causal: Boolean, window: Int, hasLengths: Boolean): String { + val positional = causal || window > 0 + return when { + hasLengths && positional -> KIND_COMBINED + hasLengths -> KIND_KEY_PADDING + window > 0 -> KIND_SLIDING_WINDOW + causal -> KIND_CAUSAL + else -> KIND_NONE + } + } + + /** Whether a node with this `maskKind` carries a `keyLengths` operand. */ + public fun hasLengths(maskKind: String?): Boolean = maskKind == KIND_KEY_PADDING || maskKind == KIND_COMBINED + + /** Argument contract shared by every implementor of the op. */ + public fun validate(window: Int, rightContext: Int) { + require(window >= 0) { "SDPA: window must be >= 0, got $window" } + require(rightContext >= 0) { "SDPA: rightContext must be >= 0, got $rightContext" } + require(rightContext == 0 || window > 0) { "SDPA: rightContext ($rightContext) is only meaningful with window > 0" } + } + + /** Absolute position of query row [qi]: the last query aligns with the last key. */ + public fun queryPosition(qi: Int, seqQ: Int, seqKV: Int): Int = seqKV - seqQ + qi + + /** First attended key (inclusive) of query row [qi]. */ + public fun firstKey(qi: Int, seqQ: Int, seqKV: Int, window: Int): Int = + if (window > 0) maxOf(0, queryPosition(qi, seqQ, seqKV) - window + 1) else 0 + + /** + * End of the attended keys (exclusive) of query row [qi]; [length] is the row's key-padding + * length (`seqKV` when there is none). The result may be `<=` [firstKey]: an empty range. + */ + public fun endKey(qi: Int, seqQ: Int, seqKV: Int, causal: Boolean, window: Int, rightContext: Int, length: Int): Int { + var end = minOf(seqKV, length) + val absQ = queryPosition(qi, seqQ, seqKV) + if (causal) end = minOf(end, absQ + 1) + if (window > 0) end = minOf(end, absQ + rightContext + 1) + return maxOf(0, end) + } + + /** The attended keys of query row [qi] as a half-open range `[first, end)`, possibly empty. */ + public fun keyRange(qi: Int, seqQ: Int, seqKV: Int, causal: Boolean, window: Int, rightContext: Int, length: Int): IntRange = + firstKey(qi, seqQ, seqKV, window) until endKey(qi, seqQ, seqKV, causal, window, rightContext, length) + + /** Whether query row [qi] attends to key [ki] under the given rules. */ + public fun attends(qi: Int, ki: Int, seqQ: Int, seqKV: Int, causal: Boolean, window: Int, rightContext: Int, length: Int): Boolean = + ki in keyRange(qi, seqQ, seqKV, causal, window, rightContext, length) + + /** + * The dense additive form of the rules: `[batch, 1, seqQ, seqKV]` row-major, `0` where attended + * and [MASKED] elsewhere — the tensor a caller built by hand before the op took the rules as + * metadata, and the reference the structured paths are tested against bit for bit. + * + * @param lengths per-row key-padding lengths (`size == batch`), or `null` for no padding + */ + public fun additive(batch: Int, seqQ: Int, seqKV: Int, causal: Boolean, window: Int, rightContext: Int, lengths: IntArray?): FloatArray { + validate(window, rightContext) + require(lengths == null || lengths.size == batch) { "SDPA: keyLengths has ${lengths?.size} rows, batch is $batch" } + val out = FloatArray(batch * seqQ * seqKV) { MASKED } + for (b in 0 until batch) { + val length = lengths?.get(b) ?: seqKV + for (qi in 0 until seqQ) { + val row = (b * seqQ + qi) * seqKV + for (ki in keyRange(qi, seqQ, seqKV, causal, window, rightContext, length)) out[row + ki] = 0f + } + } + return out + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt index 440e63cf7..6cab2e6b2 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt @@ -1014,11 +1014,44 @@ public class UnsqueezeOperation( /** * Scaled dot-product attention operation for tape recording. * Output shape = query shape: [batch, nHeads, seqLen, headDim] + * + * Operands are `query, key, value` followed by the optional dense `mask` and the optional + * `keyLengths` (#1343), in that order and only when present; the parameters carry `scale`, + * `causal`, `window`, `rightContext` and `maskKind` ([AttentionMask.kindOf]). [operandLayout] + * resolves which trailing operand is which, so every consumer reads the same rule. */ public class ScaledDotProductAttentionOperation( parameters: Map = emptyMap() ) : BaseOperation("scaledDotProductAttention", "nn", parameters) { + /** Index of the dense mask operand and of the `keyLengths` operand, each `null` when absent. */ + public data class OperandLayout(val maskIndex: Int?, val lengthsIndex: Int?) + + public companion object { + /** + * Resolve the trailing operands of an SDPA node with [inputCount] operands: `keyLengths` + * is the last operand when `parameters["maskKind"]` says the node carries lengths + * ([AttentionMask.hasLengths]); the dense mask is operand 3 when one more operand remains. + * Nodes recorded before `maskKind` existed have no lengths operand. + */ + public fun operandLayout(inputCount: Int, parameters: Map): OperandLayout { + val hasLengths = AttentionMask.hasLengths(parameters["maskKind"] as? String) + val lengthsIndex = if (hasLengths && inputCount >= 4) inputCount - 1 else null + val tensorCount = if (lengthsIndex != null) inputCount - 1 else inputCount + val maskIndex = if (tensorCount >= 4) 3 else null + return OperandLayout(maskIndex, lengthsIndex) + } + + /** The parameters an SDPA call records, with `maskKind` derived from the rules in force. */ + public fun parameters(scale: Float, causal: Boolean, window: Int, rightContext: Int, hasLengths: Boolean): Map = mapOf( + "scale" to scale, + "causal" to causal, + "window" to window, + "rightContext" to rightContext, + "maskKind" to AttentionMask.kindOf(causal, window, hasLengths), + ) + } + override fun execute( inputs: List> ): List> = emptyList() diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index 0e82b8750..fffb31dca 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -394,21 +394,40 @@ public interface TensorOps { * attends to K/V head `h / (nHeads / nKVHeads)` — the same mapping the transformer modules * use — so K/V are never tiled to the query head count (SKEEP-005 phase 2). * + * Masks are *rules over positions*, defined once in [AttentionMask] and shared by the CPU + * kernel, the recorded graph and the StableHLO lowering (#1343): [causal], a sliding + * [window] with [rightContext] lookahead, and key padding given as per-row [keyLengths]. + * Several rules intersect; a masked score gets [AttentionMask.MASKED] (`-1e30`, finite). + * The rules are recorded as `parameters["causal" | "window" | "rightContext" | "maskKind"]` + * and [keyLengths] as the last operand, so a compiled target can lower them to a fused kernel + * that takes lengths or to an `iota`/`compare`/`select` decomposition instead of inheriting a + * `[batch, 1, seqLen, kvLen]` constant. The dense additive [mask] remains the generic + * fallback and is added on top of the rules. + * * @param query [batch, nHeads, seqLen, headDim] * @param key [batch, nKVHeads, kvLen, headDim] * @param value [batch, nKVHeads, kvLen, headDim] - * @param mask optional additive mask [batch, 1, seqLen, kvLen] (e.g. causal) - * @param scale scaling factor, defaults to 1/sqrt(headDim) - * @param causal if true, apply causal masking (ignore [mask] parameter) + * @param mask optional additive mask broadcastable to [batch, 1, seqLen, kvLen] + * @param scale scaling factor, `0` selects 1/sqrt(headDim) + * @param causal if true, query `q` attends to keys `k <= kvLen - seqLen + q` only + * @param keyLengths optional `[batch]` integer tensor: row `b` attends to keys `< keyLengths[b]` + * @param window sliding-window size in keys, `0` for none: query `q` attends to keys + * `>= kvLen - seqLen + q - window + 1` + * @param rightContext bounded lookahead of the window band in keys (needs `window > 0`): + * query `q` attends to keys `<= kvLen - seqLen + q + rightContext` * @return [batch, nHeads, seqLen, headDim] */ @Diff + @Lowering(composite = "skainet.sdpa") public fun scaledDotProductAttention( query: Tensor, key: Tensor, value: Tensor, mask: Tensor? = null, scale: Float = 0f, - causal: Boolean = false + causal: Boolean = false, + keyLengths: Tensor? = null, + window: Int = 0, + rightContext: Int = 0 ): Tensor } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt index 973a8c8b7..8d3a764ba 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt @@ -566,7 +566,10 @@ public class VoidTensorOps : TensorOps { value: Tensor, mask: Tensor?, scale: Float, - causal: Boolean + causal: Boolean, + keyLengths: Tensor?, + window: Int, + rightContext: Int ): Tensor { // Output shape matches query shape: [batch, nHeads, seqLen, headDim] val resultData = dataFactory.zeros(query.shape, query.dtype) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/OpAttributeFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/OpAttributeFactory.kt index 461eb7cc0..a07ead138 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/OpAttributeFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/OpAttributeFactory.kt @@ -18,6 +18,14 @@ public object OpAttributeFactory { input: Tensor<*, *>, result: Tensor<*, *> ): Map = shapesAndDTypes(listOf(input), listOf(result)) + /** + * `maskKind` of a `scaledDotProductAttention` call (#1343): which positional / padding rules the + * call applies, derived from its arguments so the trace carries it beside `causal`, `window`, + * `rightContext` and `scale`. See [sk.ainet.lang.tensor.ops.AttentionMask.kindOf]. + */ + public fun sdpaMaskKind(keyLengths: Tensor<*, *>?, causal: Boolean, window: Int): Map = + mapOf("maskKind" to sk.ainet.lang.tensor.ops.AttentionMask.kindOf(causal, window, keyLengths != null)) + /** Conv2d op attributes: shapes/dtypes + stride/padding/dilation/groups and bias flag. */ public fun conv2d( input: Tensor<*, *>, diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt new file mode 100644 index 000000000..835b4480c --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt @@ -0,0 +1,77 @@ +package sk.ainet.lang.tensor.ops + +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** #1343: the single definition of the attention mask rules. */ +class AttentionMaskTest { + + @Test + fun kind_names_the_rules_in_force() { + assertEquals("none", AttentionMask.kindOf(causal = false, window = 0, hasLengths = false)) + assertEquals("causal", AttentionMask.kindOf(causal = true, window = 0, hasLengths = false)) + assertEquals("slidingWindow", AttentionMask.kindOf(causal = false, window = 4, hasLengths = false)) + assertEquals("slidingWindow", AttentionMask.kindOf(causal = true, window = 4, hasLengths = false)) + assertEquals("keyPadding", AttentionMask.kindOf(causal = false, window = 0, hasLengths = true)) + assertEquals("combined", AttentionMask.kindOf(causal = true, window = 0, hasLengths = true)) + assertEquals("combined", AttentionMask.kindOf(causal = false, window = 4, hasLengths = true)) + assertTrue(AttentionMask.hasLengths("keyPadding") && AttentionMask.hasLengths("combined")) + assertTrue(!AttentionMask.hasLengths("causal") && !AttentionMask.hasLengths(null)) + } + + @Test + fun causal_aligns_the_last_query_with_the_last_key() { + // Prefill: diagonal. + assertEquals(0 until 3, AttentionMask.keyRange(2, 4, 4, causal = true, window = 0, rightContext = 0, length = 4)) + // KV-cache decode: one query, the whole cache. + assertEquals(0 until 9, AttentionMask.keyRange(0, 1, 9, causal = true, window = 0, rightContext = 0, length = 9)) + // Chunked prefill: 2 queries over 6 keys attend up to positions 4 and 5. + assertEquals(0 until 5, AttentionMask.keyRange(0, 2, 6, causal = true, window = 0, rightContext = 0, length = 6)) + assertEquals(0 until 6, AttentionMask.keyRange(1, 2, 6, causal = true, window = 0, rightContext = 0, length = 6)) + } + + @Test + fun window_band_matches_the_transformer_formula() { + // transformers MultiHeadAttention.buildSlidingWindowMask: absQ = seqKV - seqQ + qi, + // allowed ki in [absQ - window + 1, absQ + rightContext]. + for ((seqQ, seqKV) in listOf(8 to 8, 1 to 20, 5 to 12)) for (window in listOf(1, 3, 16)) for (rightContext in listOf(0, 2)) { + val got = AttentionMask.additive(1, seqQ, seqKV, causal = false, window = window, rightContext = rightContext, lengths = null) + val expected = FloatArray(seqQ * seqKV) + for (qi in 0 until seqQ) { + val absQ = seqKV - seqQ + qi + for (ki in 0 until seqKV) expected[qi * seqKV + ki] = if (ki in (absQ - window + 1)..(absQ + rightContext)) 0f else -1.0e30f + } + assertContentEquals(expected, got, "seqQ=$seqQ seqKV=$seqKV window=$window rightContext=$rightContext") + } + } + + @Test + fun padding_intersects_with_the_positional_rules() { + val lengths = intArrayOf(3, 0, 5) + val m = AttentionMask.additive(3, 4, 5, causal = true, window = 0, rightContext = 0, lengths = lengths) + fun attended(b: Int, qi: Int, ki: Int) = m[(b * 4 + qi) * 5 + ki] == 0f + // Row 0: keys < 3 and ki <= absQ (absQ = 1 + qi). + assertTrue(attended(0, 0, 0) && attended(0, 0, 1) && !attended(0, 0, 2)) + assertTrue(attended(0, 3, 2) && !attended(0, 3, 3)) + // Row 1: fully padded → every key masked, but finite. + for (qi in 0 until 4) for (ki in 0 until 5) assertEquals(-1.0e30f, m[(1 * 4 + qi) * 5 + ki]) + // Row 2: no padding → plain causal with offset 1. + assertTrue(attended(2, 0, 1) && !attended(2, 0, 2) && attended(2, 3, 4)) + // keyRange agrees with attends, including an empty range. + assertTrue(AttentionMask.keyRange(0, 4, 5, true, 0, 0, 0).isEmpty()) + for (qi in 0 until 4) for (ki in 0 until 5) { + assertEquals(attended(0, qi, ki), AttentionMask.attends(qi, ki, 4, 5, causal = true, window = 0, rightContext = 0, length = 3)) + } + } + + @Test + fun argument_contract() { + assertFailsWith { AttentionMask.validate(window = -1, rightContext = 0) } + assertFailsWith { AttentionMask.validate(window = 0, rightContext = 1) } + assertFailsWith { AttentionMask.additive(2, 2, 2, false, 0, 0, intArrayOf(1)) } + AttentionMask.validate(window = 4, rightContext = 2) + } +} diff --git a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api index 3bd72d832..5f37b32cc 100644 --- a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api +++ b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api @@ -127,8 +127,8 @@ public final class sk/ainet/lang/dag/GraphDslOpsGraphDslKt { public static synthetic fun reshape$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/tensor/Shape;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun rsubScalar (Lsk/ainet/lang/dag/DagBuilder;Ljava/lang/Number;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun rsubScalar$default (Lsk/ainet/lang/dag/DagBuilder;Ljava/lang/Number;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; - public static final fun scaledDotProductAttention (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;FZLjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; - public static synthetic fun scaledDotProductAttention$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;FZLjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; + public static final fun scaledDotProductAttention (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;FZLsk/ainet/lang/dag/GraphValue;IILjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; + public static synthetic fun scaledDotProductAttention$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;FZLsk/ainet/lang/dag/GraphValue;IILjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun sigmoid (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun sigmoid$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun sign (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt index b6748061a..30b824639 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt @@ -73,7 +73,10 @@ class GraphDslGenerator( val dslParams = method.parameters.map { param -> val dslType = if (param.isTensor) { - if (param.type.endsWith("?")) "GraphValue<*>?" else "GraphValue<*>" + // An optional tensor parameter of TensorOps always defaults to null (mask, bias, + // keyLengths); KSP's rendered type may have lost the `?`, so treat it as nullable + // and give the DSL the same `= null` default instead of a mandatory positional. + if (param.type.endsWith("?") || param.isOptional) "GraphValue<*>?" else "GraphValue<*>" } else if (param.type.contains("Tensor<")) { if (param.type.startsWith("kotlin.collections.List<") || param.type.startsWith("List<")) "List>" else "GraphValue<*>" } else { @@ -100,6 +103,7 @@ class GraphDslGenerator( else if (param.name == "dim") " = -1" else if (param.name == "startDim") " = 0" else if (param.name == "endDim") " = -1" + else if (param.name == "window" || param.name == "rightContext") " = 0" else "" } else if (param.type == "sk.ainet.lang.tensor.ops.UpsampleMode") { " = UpsampleMode.Nearest" diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt index 59d9b1244..f3b4c4b81 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt @@ -326,6 +326,12 @@ class MethodAnalyzer { generateConv2dCall(method) } + // Scaled dot-product attention (#1343): the trace carries `maskKind`, derived from the + // rules in force, beside the plain scalar attributes (scale, causal, window, rightContext). + method.name == "scaledDotProductAttention" && method.parameters.any { it.name == "keyLengths" } -> { + "sdpaMaskKind(keyLengths as Tensor<*, *>?, causal, window)" + } + // Shape operations - use shapesAndDTypes method.name in listOf("reshape", "flatten", "squeeze", "unsqueeze", "transpose", "concat", "split") -> { val inputs = tensorParams.map { "${it.name} as Tensor<*, *>" }.joinToString(", ") diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt index 0f4306f1a..1784d2448 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt @@ -336,7 +336,7 @@ class TracingWrapperGenerator( val correctedType = when { param.name == "dim" && method.name in listOf("squeeze", "sum", "mean", "variance") && !param.type.endsWith("?") -> "${param.type}?" param.name == "bias" && method.name in listOf("conv1d", "conv2d", "conv3d", "convTranspose1d") && !param.type.endsWith("?") -> "${param.type}?" - param.name == "mask" && method.name == "scaledDotProductAttention" && !param.type.endsWith("?") -> "${param.type}?" + param.name in listOf("mask", "keyLengths") && method.name == "scaledDotProductAttention" && !param.type.endsWith("?") -> "${param.type}?" else -> param.type } diff --git a/skainet-lang/skainet-lang-ksp-processor/src/test/resources/sk/ainet/lang/trace/OpAttributeFactory.kt b/skainet-lang/skainet-lang-ksp-processor/src/test/resources/sk/ainet/lang/trace/OpAttributeFactory.kt index 257a7b318..a8ebe2d45 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/test/resources/sk/ainet/lang/trace/OpAttributeFactory.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/test/resources/sk/ainet/lang/trace/OpAttributeFactory.kt @@ -20,6 +20,9 @@ object OpAttributeFactory { ) } + fun sdpaMaskKind(keyLengths: Any?, causal: Boolean, window: Int): Map = + mapOf("maskKind" to if (keyLengths != null) "keyPadding" else if (causal) "causal" else "none") + fun conv2d( input: Any, weight: Any, From 89b052f10e9ca2edac7c50d952276d3ef828fb2a Mon Sep 17 00:00:00 2001 From: michalharakal Date: Thu, 15 Oct 2026 17:19:21 +0200 Subject: [PATCH 2/2] test(lang-core): compare masked mask entries by threshold on Kotlin/JS Kotlin/JS keeps a Float literal at double precision while a FloatArray element is float32, so the -1e30 equality in AttentionMaskTest failed on the JS leg. --- .../kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt index 835b4480c..0e0c735ee 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/AttentionMaskTest.kt @@ -56,8 +56,12 @@ class AttentionMaskTest { // Row 0: keys < 3 and ki <= absQ (absQ = 1 + qi). assertTrue(attended(0, 0, 0) && attended(0, 0, 1) && !attended(0, 0, 2)) assertTrue(attended(0, 3, 2) && !attended(0, 3, 3)) - // Row 1: fully padded → every key masked, but finite. - for (qi in 0 until 4) for (ki in 0 until 5) assertEquals(-1.0e30f, m[(1 * 4 + qi) * 5 + ki]) + // Row 1: fully padded → every key masked, but finite. (Threshold, not equality: Kotlin/JS + // keeps a Float literal at double precision while a FloatArray element is float32.) + for (qi in 0 until 4) for (ki in 0 until 5) { + val v = m[(1 * 4 + qi) * 5 + ki] + assertTrue(v.isFinite() && v <= -1.0e29f, "row 1 key $ki of query $qi must be masked, got $v") + } // Row 2: no padding → plain causal with offset 1. assertTrue(attended(2, 0, 1) && !attended(2, 0, 2) && attended(2, 3, 4)) // keyRange agrees with attends, including an empty range.