[michalharakal]

Summary

Part 1 of #1343: key padding, the sliding-window band and causal become rules over positions that scaledDotProductAttention carries as metadata, with one definition shared by the CPU kernel, the recorded graph and the StableHLO lowering.

• TensorOps.scaledDotProductAttention(…, keyLengths: Tensor<DType, *>? = null, window: Int = 0, rightContext: Int = 0) (@Lowering(composite = "skainet.sdpa")). keyLengths is a [batch] integer tensor (row b attends to keys < keyLengths[b]); the window band is MultiHeadAttention's formula (absQ = Sk − Sq + qi, keys in [absQ − window + 1, absQ + rightContext]). Rules intersect; the dense additive mask stays the generic fallback and is added on top.
• AttentionMask (sk.ainet.lang.tensor.ops): kindOf (none | causal | slidingWindow | keyPadding | combined), keyRange / attends / endKey, additive (the [b, 1, Sq, Sk] tensor callers built by hand, now the reference the structured paths are tested against), MASKED = -1e30 (finite, never -inf).
• CPU kernel: when lengths or a window are given, scores are computed only inside each row's key range and the rest is MASKED — bit-identical to the dense -1e30 path on the attended keys, uniform (finite) on a fully padded row, causal folded into the range. Without them the kernel is untouched byte for byte.
• Recording: RecordingExecution and the KSP tracer (OpAttributeFactory.sdpaMaskKind) record maskKind, window, rightContext beside scale/causal, with keyLengths as the last operand after the optional dense mask; ScaledDotProductAttentionOperation.operandLayout resolves the layout for ComputeGraphExecutor and the converter (legacy nodes without maskKind keep their meaning).
• StableHLO: the converter body moved into SdpaLowering.emit(e: HloEmitter, operands, operandTypes, resultType, params), which reads every shape from the operand MLIR types — the contract a stablehlo.composite decomposition function has (part 2 wires skainet.sdpa onto the n-ary seam). Rules lower to one i1 keep mask (iota/compare/and) and select 0 / -1e30 added to the scores; keyLengths is a plain [batch] i32 input, no [b, 1, Sq, Sk] constant enters the module.
• GraphDsl (skainet-lang-dag): optional tensor operands (mask, bias, keyLengths) are now GraphValue<*>? = null instead of mandatory positionals.

Behaviour changes on the export side (flagged in the changelog)

• The causal diagonal carries the Sk − Sq offset the CPU kernel uses, so a KV-cache decode step attends to its whole cache. Sq == Sk emission is unchanged (the pre-change SdpaNumericDumpTest dump is reproduced line for line; only the SSA numbering order differs because the builder names values lazily).
• A dense mask and causal = true apply together, as on the CPU, instead of the mask silently replacing the rule.
• Rules on a dynamic key length (?) return a conversion failure instead of invalid iota IR; a dense mask still exports on the dynamic path.

Out of scope, stated

• scaledDotProductAttentionBackward stays the existing stub (null gradients); when a backward is written it takes its mask from AttentionMask like the two forward paths.
• stablehlo-legalize-composite-to-call round trip: no stablehlo-opt on the CI/dev machines; part 2 keeps the "composite body ≡ inline modulo SSA names" test instead.

Test plan

• AttentionMaskTest (lang-core): kinds, causal offset, window band == MultiHeadAttention formula, padding ∩ rules, contract
• SdpaKeyLengthsTest (CPU): lengths == dense mask bit for bit (causal on/off, GQA, Sq ≠ Sk), window/rightContext == dense, combined, padded batch == single passes ≤ 1e-6, fully padded row finite, schedule-independent, argument errors
• SdpaKeyLengthsRecordingTest (compile-dag): trace/node metadata, operand layout for every combination, DIRECT vs OPTIMIZED ≤ 1e-6
• SdpaMaskRulesHloTest (compile-hlo): lengths → iota/compare/select from the i32 input, band edges + causal offset, Sq == Sk causal lines unchanged, dense mask + rules, dynamic key length refused / dense mask accepted
• existing SDPA suites (SdpaGqaParityTest, SdpaScheduleParityTest, SDPAShapeValidationTest, SdpaGqaHloExportTest, DynamicShapeHloExportTest, ScheduleModuleAttributeTest, NarrowPermuteConverterTest) unchanged
• apiDump, generateDocs (matrix row), scripts/pr-gate.sh --quick, full gate