From 8e0de7841de11f26c7d5841a063e707587a418af Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 13:54:27 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20sliding-window=20KV=20=E2=80=94?= =?UTF-8?q?=20a=20wrapped=20window=20is=20a=20pair=20of=20views,=20and=20a?= =?UTF-8?q?ttention=20iterates=20it?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1036 (SKEEP-003 P4, S2.3, proposal §4.6 ring decision, decision #4; M2-F5, M2-A4). A ring that has wrapped holds its newest positions in two runs, and the choice has always been to copy them together before attention or to grow the buffer forever. This is the third option: hand the kernel the pair. - `WindowedKV(head, tail?)` — the window as the one or two runs it physically is. Both halves are ordinary TensorViews over the same Storage, so nothing is copied; `gather(scope, sink)` makes the copy *visible* as one traced adapter for kernels that cannot iterate a pair. - `KvCacheStore.keyWindow/valueWindow` are defaulted (the default copies once through readKeys, so a TurboQuant store needs no special case) and overridden zero-copy by `DefaultKvCacheStore`: strided views with the head stride left at `maxSeqLen * headDim`, and two of them when the run crosses the end of the ring. - `DefaultKvCacheStore(slidingWindow = true)` makes the buffer a ring: appending past capacity overwrites the oldest position, positions stay absolute, `windowStart` is the oldest one still held, and reads across the wrap return positions in order. Off by default — a full cache still throws exactly as before. - `WindowedAttention.decodeStep` computes softmax online (running maximum and denominator), which is what makes the pair workable: positions are consumed run after run and nothing has to exist as one array. The accumulator is the caller's output row, so it allocates nothing per token. `CompressedKvAttention` exposes the same windows. M2-A4 is asserted directly: a ring that wrapped several times and a cache that never wrapped, over the same four positions, produce bit-identical output (`assertContentEquals`) — and the test first asserts the window really is wrapped, so it cannot pass vacuously. The gather path agrees within 1e-5 and emits exactly two adapters priced in the trace; sixteen decode steps over a wrapped ring emit zero Allocation and zero AdapterInserted events. Gate: scripts/pr-gate.sh — all legs passed. Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-lang-core.api | 50 ++++- .../sk/ainet/lang/memory/WindowedAttention.kt | 97 +++++++++ .../kotlin/sk/ainet/lang/memory/WindowedKV.kt | 109 ++++++++++ .../tensor/storage/CompressedKvAttention.kt | 23 +++ .../tensor/storage/DefaultKvCacheStore.kt | 104 +++++++++- .../ainet/lang/tensor/storage/KvCacheStore.kt | 31 +++ .../sk/ainet/lang/memory/WindowedKvTest.kt | 195 ++++++++++++++++++ 7 files changed, 601 insertions(+), 8 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedAttention.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedKV.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/WindowedKvTest.kt 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 7297a5242..ec6247077 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 @@ -1246,6 +1246,35 @@ public final class sk/ainet/lang/memory/ViewsKt { public static final fun viewOrNull (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/memory/TensorView; } +public final class sk/ainet/lang/memory/WindowedAttention { + public static final field INSTANCE Lsk/ainet/lang/memory/WindowedAttention; + public final fun decodeStep ([FLsk/ainet/lang/memory/WindowedKV;Lsk/ainet/lang/memory/WindowedKV;[FF)V + public static synthetic fun decodeStep$default (Lsk/ainet/lang/memory/WindowedAttention;[FLsk/ainet/lang/memory/WindowedKV;Lsk/ainet/lang/memory/WindowedKV;[FFILjava/lang/Object;)V + public final fun decodeStepGathered ([FLsk/ainet/lang/memory/WindowedKV;Lsk/ainet/lang/memory/WindowedKV;[FLsk/ainet/lang/memory/Scope;Lsk/ainet/lang/memory/trace/TraceSink;F)V + public static synthetic fun decodeStepGathered$default (Lsk/ainet/lang/memory/WindowedAttention;[FLsk/ainet/lang/memory/WindowedKV;Lsk/ainet/lang/memory/WindowedKV;[FLsk/ainet/lang/memory/Scope;Lsk/ainet/lang/memory/trace/TraceSink;FILjava/lang/Object;)V +} + +public final class sk/ainet/lang/memory/WindowedKV { + public static final field Companion Lsk/ainet/lang/memory/WindowedKV$Companion; + public static final field POSITION_AXIS I + public fun (Lsk/ainet/lang/memory/TensorView;Lsk/ainet/lang/memory/TensorView;)V + public synthetic fun (Lsk/ainet/lang/memory/TensorView;Lsk/ainet/lang/memory/TensorView;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun gather (Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/memory/trace/TraceSink;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun gather$default (Lsk/ainet/lang/memory/WindowedKV;Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/memory/trace/TraceSink;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; + public final fun get (III)F + public final fun getHead ()Lsk/ainet/lang/memory/TensorView; + public final fun getHeadDim ()I + public final fun getHeads ()I + public final fun getLength ()I + public final fun getParts ()Ljava/util/List; + public final fun getTail ()Lsk/ainet/lang/memory/TensorView; + public final fun getWrapped ()Z + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/WindowedKV$Companion { +} + public final class sk/ainet/lang/memory/plan/ActualMemory { public static final field Companion Lsk/ainet/lang/memory/plan/ActualMemory$Companion; public fun (Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;)V @@ -6625,6 +6654,8 @@ public final class sk/ainet/lang/tensor/storage/CompressedKvAttention { public synthetic fun (Lsk/ainet/lang/tensor/storage/KvCacheStore;Lsk/ainet/lang/tensor/storage/CompressedKvAttention$DequantStrategy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun isKeyCompressed ()Z public final fun isValueCompressed ()Z + public final fun keyWindowForAttention (III)Lsk/ainet/lang/memory/WindowedKV; + public static synthetic fun keyWindowForAttention$default (Lsk/ainet/lang/tensor/storage/CompressedKvAttention;IIIILjava/lang/Object;)Lsk/ainet/lang/memory/WindowedKV; public final fun loadKeyStorageRaw (III)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun loadKeyStorageRaw$default (Lsk/ainet/lang/tensor/storage/CompressedKvAttention;IIIILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun loadKeysForAttention (III)[F @@ -6634,6 +6665,8 @@ public final class sk/ainet/lang/tensor/storage/CompressedKvAttention { public final fun loadValuesForAttention (III)[F public static synthetic fun loadValuesForAttention$default (Lsk/ainet/lang/tensor/storage/CompressedKvAttention;IIIILjava/lang/Object;)[F public final fun storeKeyValue (I[F[F)V + public final fun valueWindowForAttention (III)Lsk/ainet/lang/memory/WindowedKV; + public static synthetic fun valueWindowForAttention$default (Lsk/ainet/lang/tensor/storage/CompressedKvAttention;IIIILjava/lang/Object;)Lsk/ainet/lang/memory/WindowedKV; } public final class sk/ainet/lang/tensor/storage/CompressedKvAttention$DequantStrategy : java/lang/Enum { @@ -6671,7 +6704,8 @@ public final class sk/ainet/lang/tensor/storage/DefaultBufferResolver : sk/ainet public final class sk/ainet/lang/tensor/storage/DefaultKvCacheStore : sk/ainet/lang/tensor/storage/KvCacheStore { public fun (Lsk/ainet/lang/tensor/storage/KvCacheConfig;)V public fun (Lsk/ainet/lang/tensor/storage/KvCacheConfig;Lsk/ainet/lang/memory/ModelScope;)V - public synthetic fun (Lsk/ainet/lang/tensor/storage/KvCacheConfig;Lsk/ainet/lang/memory/ModelScope;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/lang/tensor/storage/KvCacheConfig;Lsk/ainet/lang/memory/ModelScope;Z)V + public synthetic fun (Lsk/ainet/lang/tensor/storage/KvCacheConfig;Lsk/ainet/lang/memory/ModelScope;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun appendToken (I[F[F)V public fun clear ()V public fun evict (I)V @@ -6685,14 +6719,18 @@ public final class sk/ainet/lang/tensor/storage/DefaultKvCacheStore : sk/ainet/l public fun getNumLayers ()I public fun getPlacement ()Lsk/ainet/lang/tensor/storage/Placement; public final fun getPreallocatedBytes ()J + public final fun getSlidingWindow ()Z public fun getValueBytesPerElement ()D public fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getValueFormat ()Lsk/ainet/lang/memory/Format; + public final fun getWindowStart ()I + public fun keyWindow (III)Lsk/ainet/lang/memory/WindowedKV; public fun memoryReport ()Lsk/ainet/lang/tensor/storage/KvCacheMemoryReport; public fun readKeyStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage; public fun readKeys (III)[F public fun readValueStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage; public fun readValues (III)[F + public fun valueWindow (III)Lsk/ainet/lang/memory/WindowedKV; } public final class sk/ainet/lang/tensor/storage/DeviceKind : java/lang/Enum { @@ -6811,6 +6849,8 @@ public abstract interface class sk/ainet/lang/tensor/storage/KvCacheStore { public fun getValueBytesPerElement ()D public abstract fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getValueFormat ()Lsk/ainet/lang/memory/Format; + public fun keyWindow (III)Lsk/ainet/lang/memory/WindowedKV; + public static synthetic fun keyWindow$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/memory/WindowedKV; public abstract fun memoryReport ()Lsk/ainet/lang/tensor/storage/KvCacheMemoryReport; public abstract fun readKeyStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun readKeyStorage$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; @@ -6820,6 +6860,8 @@ public abstract interface class sk/ainet/lang/tensor/storage/KvCacheStore { public static synthetic fun readValueStorage$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; public abstract fun readValues (III)[F public static synthetic fun readValues$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)[F + public fun valueWindow (III)Lsk/ainet/lang/memory/WindowedKV; + public static synthetic fun valueWindow$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/memory/WindowedKV; } public final class sk/ainet/lang/tensor/storage/KvCacheStore$Companion { @@ -6835,10 +6877,14 @@ public final class sk/ainet/lang/tensor/storage/KvCacheStore$DefaultImpls { public static fun getKeyFormat (Lsk/ainet/lang/tensor/storage/KvCacheStore;)Lsk/ainet/lang/memory/Format; public static fun getValueBytesPerElement (Lsk/ainet/lang/tensor/storage/KvCacheStore;)D public static fun getValueFormat (Lsk/ainet/lang/tensor/storage/KvCacheStore;)Lsk/ainet/lang/memory/Format; + public static fun keyWindow (Lsk/ainet/lang/tensor/storage/KvCacheStore;III)Lsk/ainet/lang/memory/WindowedKV; + public static synthetic fun keyWindow$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/memory/WindowedKV; public static synthetic fun readKeyStorage$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun readKeys$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)[F public static synthetic fun readValueStorage$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun readValues$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)[F + public static fun valueWindow (Lsk/ainet/lang/tensor/storage/KvCacheStore;III)Lsk/ainet/lang/memory/WindowedKV; + public static synthetic fun valueWindow$default (Lsk/ainet/lang/tensor/storage/KvCacheStore;IIIILjava/lang/Object;)Lsk/ainet/lang/memory/WindowedKV; } public final class sk/ainet/lang/tensor/storage/LogicalDType : java/lang/Enum { @@ -7350,11 +7396,13 @@ public final class sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore : sk/aine public fun getValueBytesPerElement ()D public fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getValueFormat ()Lsk/ainet/lang/memory/Format; + public fun keyWindow (III)Lsk/ainet/lang/memory/WindowedKV; public fun memoryReport ()Lsk/ainet/lang/tensor/storage/KvCacheMemoryReport; public fun readKeyStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage; public fun readKeys (III)[F public fun readValueStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage; public fun readValues (III)[F + public fun valueWindow (III)Lsk/ainet/lang/memory/WindowedKV; } public abstract interface annotation class sk/ainet/lang/tensor/storage/Weights : java/lang/annotation/Annotation { diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedAttention.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedAttention.kt new file mode 100644 index 000000000..6769fa591 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedAttention.kt @@ -0,0 +1,97 @@ +package sk.ainet.lang.memory + +import kotlin.math.exp +import kotlin.math.max +import kotlin.math.sqrt + +/** + * Attention over a [WindowedKV] — the reference kernel that **iterates the pair** instead of + * requiring one contiguous window (SKEEP-003 §4.6, M2-F5/M2-A4). + * + * The softmax is computed online (running maximum and running denominator, the flash-attention + * recurrence), which is what makes the pair workable: positions are consumed in order, one run + * after the other, and nothing has to exist as a single array. It also means the kernel allocates + * **nothing** per token — the running accumulator is the caller's output row. + * + * A kernel that cannot do this calls [WindowedKV.gather] instead and gets one contiguous view plus + * one traced adapter; both paths are asserted to agree. + */ +@ExperimentalMemoryApi +public object WindowedAttention { + + /** + * One decode step: `softmax(q·kᵀ · scale) @ v` for every head, over the window's positions in + * order. + * + * @param query `[heads * headDim]` — the current token's query rows + * @param keys the key window; `[heads, positions, headDim]` across its runs + * @param values the value window, the same shape as [keys] + * @param out `[heads * headDim]`, overwritten with the attention output + * @param scale multiplier on the scores; `0` means the usual `1/sqrt(headDim)` + */ + public fun decodeStep( + query: FloatArray, + keys: WindowedKV, + values: WindowedKV, + out: FloatArray, + scale: Float = 0f, + ) { + val heads = keys.heads + val headDim = keys.headDim + require(keys.length == values.length) { "key and value windows differ: ${keys.length} vs ${values.length}" } + require(values.heads == heads && values.headDim == headDim) { "key and value windows must have the same geometry" } + require(query.size == heads * headDim) { "query must be [heads * headDim] = ${heads * headDim}, was ${query.size}" } + require(out.size == heads * headDim) { "out must be [heads * headDim] = ${heads * headDim}, was ${out.size}" } + require(keys.length > 0) { "cannot attend over an empty window" } + val s = if (scale != 0f) scale else 1f / sqrt(headDim.toFloat()) + + for (h in 0 until heads) { + val base = h * headDim + for (d in 0 until headDim) out[base + d] = 0f + var runningMax = Float.NEGATIVE_INFINITY + var denominator = 0f + var position = 0 + // Runs in order, oldest first: head, then tail when the ring wrapped. + for (partIndex in keys.parts.indices) { + val k = keys.parts[partIndex] + val v = values.parts[partIndex] + require(k.shape[WindowedKV.POSITION_AXIS] == v.shape[WindowedKV.POSITION_AXIS]) { + "key and value runs must line up (run $partIndex)" + } + for (p in 0 until k.shape[WindowedKV.POSITION_AXIS]) { + var score = 0f + for (d in 0 until headDim) score += query[base + d] * k.get(h, p, d) + score *= s + + val newMax = max(runningMax, score) + val correction = if (runningMax == Float.NEGATIVE_INFINITY) 0f else exp(runningMax - newMax) + val weight = exp(score - newMax) + denominator = denominator * correction + weight + for (d in 0 until headDim) { + out[base + d] = out[base + d] * correction + weight * v.get(h, p, d) + } + runningMax = newMax + position++ + } + } + for (d in 0 until headDim) out[base + d] = out[base + d] / denominator + } + } + + /** + * [decodeStep] through the gather adapter: the window is copied into [scope] as one contiguous + * view first, for kernels that cannot iterate a pair. Same numbers, one allocation and one + * traced adapter per call. + */ + public fun decodeStepGathered( + query: FloatArray, + keys: WindowedKV, + values: WindowedKV, + out: FloatArray, + scope: Scope, + sink: sk.ainet.lang.memory.trace.TraceSink = sk.ainet.lang.memory.trace.NoopTraceSink, + scale: Float = 0f, + ) { + decodeStep(query, WindowedKV(keys.gather(scope, sink)), WindowedKV(values.gather(scope, sink)), out, scale) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedKV.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedKV.kt new file mode 100644 index 000000000..eb090c0bb --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/WindowedKV.kt @@ -0,0 +1,109 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.TensorId + +/** + * An attention window over a KV ring, as the **one or two pieces it physically is** (SKEEP-003 + * §4.6, decision #4; M2-F5). + * + * A ring that has wrapped holds its newest `n` positions in two runs — `[from, capacity)` and + * `[0, wrapEnd)` — and the choice has always been to copy them together before attention or to + * grow the buffer forever. This is the third option: hand the kernel the pair. Both halves are + * ordinary [TensorView]s over the same [Storage], so nothing is copied and nothing is special-cased; + * a kernel that cannot iterate a pair calls [gather], which makes the copy *visible* as one + * adapter in the trace instead of an invisible per-token allocation. + * + * Position order is `head` then `tail`: oldest to newest. + * + * @property head the first (older) run — the only one when the window does not wrap + * @property tail the second (newer) run, present only after the ring wraps + */ +@ExperimentalMemoryApi +public class WindowedKV( + public val head: TensorView, + public val tail: TensorView? = null, +) { + init { + val t = tail + if (t != null) { + require(t.shape.rank == head.shape.rank) { "both halves must have the same rank" } + require(t.format == head.format) { "both halves must have the same format" } + for (axis in 0 until head.shape.rank) { + if (axis == POSITION_AXIS) continue + require(t.shape[axis] == head.shape[axis]) { "halves differ on axis $axis: ${head.shape} vs ${t.shape}" } + } + } + } + + /** The halves in position order, oldest first. */ + public val parts: List get() = if (tail == null) listOf(head) else listOf(head, tail) + + /** Positions in the window, across both halves. */ + public val length: Int get() = head.shape[POSITION_AXIS] + (tail?.shape?.get(POSITION_AXIS) ?: 0) + + /** True when the ring wrapped inside this window — the case the pair exists for. */ + public val wrapped: Boolean get() = tail != null + + /** `[heads, dim]` — the shape of the window with its position axis removed. */ + public val heads: Int get() = head.shape[0] + public val headDim: Int get() = head.shape[head.shape.rank - 1] + + /** + * The window as one contiguous view in [scope] — the gather adapter (§5.1). + * + * For a kernel that cannot iterate the pair. The copy is real, so it is *traced*: one + * [TraceEvent.AdapterInserted] per call, which is what makes "the zero-copy path allocates + * nothing per token" an assertion rather than a hope. + */ + public fun gather(scope: Scope, sink: TraceSink = NoopTraceSink, id: TensorId? = null): TensorView { + val total = length + val shape = Shape(heads, total, headDim) + val storage = scope.allocateFloats(heads * total * headDim, id) + val out = storage.floats!! + val base = storage.arrayOffset + var written = 0 + for (part in parts) { + val positions = part.shape[POSITION_AXIS] + for (h in 0 until heads) { + for (p in 0 until positions) { + val dst = base + (h.toLong() * total + written + p).toInt() * headDim + for (d in 0 until headDim) out[dst + d] = part.get(h, p, d) + } + } + written += positions + } + if (sink.isEnabled) { + sink.emit( + TraceEvent.AdapterInserted( + kind = "gather-kv-window", + from = head.format, + to = Format.dense(head.format.dtype), + bytes = heads.toLong() * total * headDim * head.format.dtype.sizeInBytes, + target = id ?: head.id, + scope = scope.kind, + ), + ) + } + return TensorView.dense(storage, shape, head.format.dtype, id ?: head.id) + } + + /** The decoded value at (`head`, absolute window position, `dim`), crossing the halves. */ + public fun get(head: Int, position: Int, dim: Int): Float { + require(position in 0 until length) { "position $position outside the window (length $length)" } + val first = this.head.shape[POSITION_AXIS] + return if (position < first) this.head.get(head, position, dim) + else tail!!.get(head, position - first, dim) + } + + override fun toString(): String = + "WindowedKV(${heads}h × $length × $headDim, ${if (wrapped) "wrapped: ${head.shape[POSITION_AXIS]}+${tail!!.shape[POSITION_AXIS]}" else "contiguous"})" + + public companion object { + /** KV windows are `[heads, positions, headDim]`. */ + public const val POSITION_AXIS: Int = 1 + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/CompressedKvAttention.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/CompressedKvAttention.kt index 21dbdddf7..66cafc898 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/CompressedKvAttention.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/CompressedKvAttention.kt @@ -85,6 +85,29 @@ public class CompressedKvAttention( } } + /** + * The key window `[startPos, endPos)` as the one or two runs the cache physically holds + * (#1036, M2-F5) — the pair `WindowedAttention.decodeStep` iterates without copying. + * + * For a compressed cache each half is an ordinary view of the decoded window, so a quantized + * KV store needs no special case here; for the dense ring the halves are views over the ring + * itself and nothing is copied at all. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + public fun keyWindowForAttention( + layer: Int, + startPos: Int = 0, + endPos: Int = cache.currentSeqLen, + ): sk.ainet.lang.memory.WindowedKV = cache.keyWindow(layer, startPos, endPos) + + /** The value window; see [keyWindowForAttention]. */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + public fun valueWindowForAttention( + layer: Int, + startPos: Int = 0, + endPos: Int = cache.currentSeqLen, + ): sk.ainet.lang.memory.WindowedKV = cache.valueWindow(layer, startPos, endPos) + /** * Load raw [TensorStorage] for keys, preserving the cache's native encoding. * diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt index 4f12a6ced..66e3ddebc 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt @@ -28,6 +28,12 @@ import sk.ainet.lang.types.FP32 public class DefaultKvCacheStore @kotlin.jvm.JvmOverloads constructor( private val config: KvCacheConfig, private val scope: sk.ainet.lang.memory.ModelScope? = null, + /** + * Treat the buffer as a **ring** (#1036, M2-F5): appending past [maxSeqLen] overwrites the + * oldest position instead of failing, and the live window is the newest [maxSeqLen] positions. + * Off by default — a cache that fills up still throws, exactly as before. + */ + public val slidingWindow: Boolean = false, ) : KvCacheStore { override val numLayers: Int get() = config.numLayers @@ -72,7 +78,7 @@ public class DefaultKvCacheStore @kotlin.jvm.JvmOverloads constructor( override fun appendToken(layer: Int, key: FloatArray, value: FloatArray) { requireLayerIndex(layer) - check(_currentSeqLen < maxSeqLen) { + check(slidingWindow || _currentSeqLen < maxSeqLen) { "KV cache is full: currentSeqLen=$_currentSeqLen, maxSeqLen=$maxSeqLen" } require(key.size == numHeads * headDim) { @@ -82,7 +88,7 @@ public class DefaultKvCacheStore @kotlin.jvm.JvmOverloads constructor( "Value size mismatch: expected ${numHeads * headDim}, got ${value.size}" } - val pos = _currentSeqLen + val pos = slotOf(_currentSeqLen) val layerKeys = keys[layer] val layerValues = values[layer] @@ -161,6 +167,80 @@ public class DefaultKvCacheStore @kotlin.jvm.JvmOverloads constructor( // --- Internal helpers --- + // --- the ring (#1036) --------------------------------------------------------------------- + + /** Physical slot of an absolute position. Identity unless this store is a ring. */ + private fun slotOf(position: Int): Int = if (slidingWindow) position % maxSeqLen else position + + /** The oldest absolute position still held — everything before it has been overwritten. */ + public val windowStart: Int + get() = if (slidingWindow) (_currentSeqLen - maxSeqLen).coerceAtLeast(0) else 0 + + /** Storage handles over the per-layer arrays, made once so a window costs no allocation. */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + private val keyStorages by lazy { List(numLayers) { sk.ainet.lang.memory.Storage.Heap.wrap(keys[it]) } } + + @sk.ainet.lang.memory.ExperimentalMemoryApi + private val valueStorages by lazy { List(numLayers) { sk.ainet.lang.memory.Storage.Heap.wrap(values[it]) } } + + @sk.ainet.lang.memory.ExperimentalMemoryApi + override fun keyWindow(layer: Int, from: Int, to: Int): sk.ainet.lang.memory.WindowedKV = + window(keyStorages[layer], layer, from, to, config.keyDType) + + @sk.ainet.lang.memory.ExperimentalMemoryApi + override fun valueWindow(layer: Int, from: Int, to: Int): sk.ainet.lang.memory.WindowedKV = + window(valueStorages[layer], layer, from, to, config.valueDType) + + /** + * `[from, to)` as one or two strided views over the layer's array — zero copies. + * + * The layer is laid out `[head, position, dim]`, so a run of positions is a view with the head + * stride left at `maxSeqLen * headDim`: contiguous per head, strided across heads. When the run + * crosses the end of the ring it becomes two such views, oldest first. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + private fun window( + storage: sk.ainet.lang.memory.Storage, + layer: Int, + from: Int, + to: Int, + dtype: sk.ainet.lang.types.DType, + ): sk.ainet.lang.memory.WindowedKV { + requireLayerIndex(layer) + require(from in windowStart..to) { "window [$from, $to) starts before the ring's oldest position ($windowStart)" } + require(to <= _currentSeqLen) { "window end $to exceeds currentSeqLen=$_currentSeqLen" } + val length = to - from + val startSlot = slotOf(from) + val firstRun = if (slidingWindow) minOf(length, maxSeqLen - startSlot) else length + val head = view(storage, startSlot, firstRun, dtype, layer, from) + val rest = length - firstRun + return if (rest <= 0) { + sk.ainet.lang.memory.WindowedKV(head) + } else { + sk.ainet.lang.memory.WindowedKV(head, view(storage, 0, rest, dtype, layer, from + firstRun)) + } + } + + @sk.ainet.lang.memory.ExperimentalMemoryApi + private fun view( + storage: sk.ainet.lang.memory.Storage, + startSlot: Int, + positions: Int, + dtype: sk.ainet.lang.types.DType, + layer: Int, + firstPosition: Int, + ): sk.ainet.lang.memory.TensorView { + val shape = Shape(numHeads, positions, headDim) + val layout = sk.ainet.lang.memory.Layout( + shape = shape, + strides = intArrayOf(maxSeqLen * headDim, headDim, 1), + offsetElements = startSlot.toLong() * headDim, + elementBytes = dtype.sizeInBytes, + ) + val id = sk.ainet.lang.tensor.TensorId(listOf("kv", "layers[$layer]"), "window", "from=$firstPosition") + return sk.ainet.lang.memory.TensorView(shape, sk.ainet.lang.memory.Format(dtype, config.keyEncoding), layout, storage, id) + } + private fun readRange( layerData: FloatArray, layer: Int, @@ -168,17 +248,27 @@ public class DefaultKvCacheStore @kotlin.jvm.JvmOverloads constructor( endPos: Int ): FloatArray { requireLayerIndex(layer) - require(startPos in 0..endPos) { "Invalid range: startPos=$startPos, endPos=$endPos" } + require(startPos in windowStart..endPos) { "Invalid range: startPos=$startPos, endPos=$endPos, oldest held position=$windowStart" } require(endPos <= _currentSeqLen) { "endPos=$endPos exceeds currentSeqLen=$_currentSeqLen" } val seqLen = endPos - startPos val result = FloatArray(numHeads * seqLen * headDim) - for (h in 0 until numHeads) { - val srcBase = h * maxSeqLen * headDim + startPos * headDim - val dstBase = h * seqLen * headDim - layerData.copyInto(result, dstBase, srcBase, srcBase + seqLen * headDim) + // A ring's window can cross the end of the buffer: copy it in the one or two runs it + // occupies, oldest first, so a wrapped read returns positions in order (#1036). + var written = 0 + var position = startPos + while (written < seqLen) { + val slot = slotOf(position) + val run = if (slidingWindow) minOf(seqLen - written, maxSeqLen - slot) else seqLen - written + for (h in 0 until numHeads) { + val srcBase = h * maxSeqLen * headDim + slot * headDim + val dstBase = h * seqLen * headDim + written * headDim + layerData.copyInto(result, dstBase, srcBase, srcBase + run * headDim) + } + written += run + position += run } return result } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt index 6a75d8e4c..748066784 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt @@ -85,6 +85,37 @@ public interface KvCacheStore { @sk.ainet.lang.memory.ExperimentalMemoryApi public val valueBytesPerElement: Double get() = kvBytesPerElement(valueFormat) + /** + * The attention window `[from, to)` of this layer's keys, as the one or two runs it physically + * occupies (#1036, M2-F5). A ring that has wrapped holds its newest positions in two runs; the + * pair is handed to attention instead of being copied together first. + * + * The default implementation copies the range once through [readKeys], which is correct for + * every store; [DefaultKvCacheStore] overrides it with zero-copy views over the ring itself. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + public fun keyWindow(layer: Int, from: Int = 0, to: Int = currentSeqLen): sk.ainet.lang.memory.WindowedKV = + copiedWindow(readKeys(layer, from, to), to - from, keyFormat.dtype) + + /** The attention window of this layer's values; see [keyWindow]. */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + public fun valueWindow(layer: Int, from: Int = 0, to: Int = currentSeqLen): sk.ainet.lang.memory.WindowedKV = + copiedWindow(readValues(layer, from, to), to - from, valueFormat.dtype) + + /** A single-run window over an already-materialized `[heads, positions, headDim]` array. */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + private fun copiedWindow( + values: FloatArray, + positions: Int, + dtype: sk.ainet.lang.types.DType, + ): sk.ainet.lang.memory.WindowedKV = sk.ainet.lang.memory.WindowedKV( + sk.ainet.lang.memory.TensorView.dense( + sk.ainet.lang.memory.Storage.Heap.wrap(values), + sk.ainet.lang.tensor.Shape(numHeads, positions, headDim), + dtype, + ), + ) + /** Placement intent for the cache buffers. */ public val placement: Placement diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/WindowedKvTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/WindowedKvTest.kt new file mode 100644 index 000000000..d16be81e7 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/WindowedKvTest.kt @@ -0,0 +1,195 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.storage.DefaultKvCacheStore +import sk.ainet.lang.tensor.storage.KvCacheConfig +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * #1036 (M2-F5, M2-A4): a sliding-window KV cache hands attention the one or two runs the ring + * physically holds, and the answer is the same as if the window had never wrapped. + */ +@OptIn(ExperimentalMemoryApi::class) +class WindowedKvTest { + + private val heads = 2 + private val headDim = 4 + + private fun store(maxSeqLen: Int, sliding: Boolean) = DefaultKvCacheStore( + KvCacheConfig(numLayers = 1, numHeads = heads, headDim = headDim, maxSeqLen = maxSeqLen), + null, + sliding, + ) + + /** Deterministic per-token K/V so a position's contents identify it. */ + private fun token(step: Int, offset: Float = 0f) = + FloatArray(heads * headDim) { i -> (step * 100 + i).toFloat() * 0.01f + offset } + + private fun fill(s: DefaultKvCacheStore, steps: Int) { + for (step in 0 until steps) s.appendToken(0, token(step), token(step, offset = 0.5f)) + } + + // --- the ring ------------------------------------------------------------------------------ + + @Test + fun aNonSlidingCacheStillRefusesToOverflow() { + val s = store(maxSeqLen = 4, sliding = false) + fill(s, 4) + assertFailsWith { s.appendToken(0, token(4), token(4)) } + assertEquals(0, s.windowStart, "without the ring, position 0 is still held") + } + + @Test + fun aRingKeepsTheNewestPositionsAndForgetsTheRest() { + val s = store(maxSeqLen = 4, sliding = true) + fill(s, 6) + assertEquals(6, s.currentSeqLen, "positions stay absolute") + assertEquals(2, s.windowStart, "the two oldest were overwritten") + assertFailsWith { s.keyWindow(0, from = 1, to = 6) } + assertFailsWith { s.readKeys(0, 0, 6) } + } + + @Test + fun aWrappedWindowIsTwoRunsInPositionOrder() { + val s = store(maxSeqLen = 4, sliding = true) + fill(s, 6) // slots hold positions 4,5,2,3 — the window 2..6 wraps + val w = s.keyWindow(0, from = 2, to = 6) + assertTrue(w.wrapped, "the window crosses the end of the ring") + assertEquals(2, w.parts.size) + assertEquals(4, w.length) + assertEquals(listOf(2, 2), w.parts.map { it.shape[WindowedKV.POSITION_AXIS] }) + + // oldest first: positions 2,3,4,5 + for ((index, step) in listOf(2, 3, 4, 5).withIndex()) { + val expected = token(step) + for (h in 0 until heads) for (d in 0 until headDim) { + assertEquals(expected[h * headDim + d], w.get(h, index, d), "position $step ($h,$d)") + } + } + } + + @Test + fun anUnwrappedWindowIsASingleRun() { + val s = store(maxSeqLen = 8, sliding = true) + fill(s, 5) + val w = s.keyWindow(0, from = 1, to = 5) + assertFalse(w.wrapped) + assertEquals(1, w.parts.size) + assertEquals(4, w.length) + } + + @Test + fun readingAcrossTheWrapReturnsPositionsInOrder() { + val s = store(maxSeqLen = 4, sliding = true) + fill(s, 7) + val flat = s.readKeys(0, 3, 7) + // [heads, 4 positions, headDim], oldest first + for ((index, step) in listOf(3, 4, 5, 6).withIndex()) { + val expected = token(step) + for (h in 0 until heads) for (d in 0 until headDim) { + assertEquals(expected[h * headDim + d], flat[(h * 4 + index) * headDim + d], "position $step") + } + } + } + + // --- the views are views ------------------------------------------------------------------- + + @Test + fun aWindowIsZeroCopyOverTheCachesOwnStorage() { + val s = store(maxSeqLen = 4, sliding = true) + fill(s, 6) + val first = s.keyWindow(0, 2, 6) + val second = s.keyWindow(0, 2, 6) + assertEquals(first.head.storage.id, second.head.storage.id, "windows share the cache's storage, they do not copy it") + assertEquals(first.head.storage.id, first.tail!!.storage.id, "and both halves are the same storage") + + // a later append is visible through a window taken before it + val w = s.keyWindow(0, 3, 6) + s.appendToken(0, token(99), token(99, 0.5f)) + assertEquals(s.keyWindow(0, 3, 7).length, w.length + 1) + } + + // --- M2-A4: the wrapped window computes what an unwrapped one would ------------------------ + + @Test + fun wrappedAttentionMatchesANonRingRunOverTheSameWindow() { + val windowLength = 4 + val ring = store(maxSeqLen = windowLength, sliding = true) + fill(ring, 10) // wrapped several times: window is 6..10 + + // the same positions, in a cache that never wrapped + val flat = store(maxSeqLen = 64, sliding = false) + for (step in 6 until 10) flat.appendToken(0, token(step), token(step, offset = 0.5f)) + + val query = FloatArray(heads * headDim) { i -> 0.1f * (i + 1) } + val fromRing = FloatArray(heads * headDim) + val fromFlat = FloatArray(heads * headDim) + WindowedAttention.decodeStep(query, ring.keyWindow(0, 6, 10), ring.valueWindow(0, 6, 10), fromRing) + WindowedAttention.decodeStep(query, flat.keyWindow(0, 0, 4), flat.valueWindow(0, 0, 4), fromFlat) + + assertTrue(ring.keyWindow(0, 6, 10).wrapped, "the ring's window must actually wrap for this to mean anything") + assertContentEquals(fromFlat, fromRing, "a wrapped window must produce the same logits as a flat one") + } + + @Test + fun theGatherAdapterAgreesWithThePairPathAndSaysSo() { + val ring = store(maxSeqLen = 4, sliding = true) + fill(ring, 6) + val keys = ring.keyWindow(0, 2, 6) + val values = ring.valueWindow(0, 2, 6) + val query = FloatArray(heads * headDim) { i -> 0.05f * (i + 3) } + + val direct = FloatArray(heads * headDim) + WindowedAttention.decodeStep(query, keys, values, direct) + + val sink = RecordingTraceSink() + val scope = ForwardScope(slabFloats = heads * 4 * headDim * 2 + 64, sink = sink, name = "attn") + val gathered = FloatArray(heads * headDim) + WindowedAttention.decodeStepGathered(query, keys, values, gathered, scope, sink) + + for (i in direct.indices) { + assertTrue(abs(direct[i] - gathered[i]) < 1e-5f, "element $i: ${direct[i]} vs ${gathered[i]}") + } + val adapters = sink.eventsOf() + assertEquals(2, adapters.size, "one visible adapter per gathered window (keys, values)") + assertTrue(adapters.all { it.kind == "gather-kv-window" }, adapters.map { it.kind }.toString()) + assertTrue(adapters.all { it.bytes == heads.toLong() * 4 * headDim * 4 }, "the copy is priced in the trace") + scope.close() + } + + @Test + fun theZeroCopyPathAllocatesNothingPerToken() { + val sink = RecordingTraceSink() + val s = store(maxSeqLen = 4, sliding = true) + fill(s, 8) + val query = FloatArray(heads * headDim) { 0.25f } + val out = FloatArray(heads * headDim) + sink.clear() + repeat(16) { + WindowedAttention.decodeStep(query, s.keyWindow(0, s.windowStart, s.currentSeqLen), s.valueWindow(0, s.windowStart, s.currentSeqLen), out) + } + assertEquals(0, sink.eventsOf().size, "windows and the pair kernel allocate nothing") + assertEquals(0, sink.eventsOf().size, "and insert no adapters") + } + + @Test + fun everyStoreCanProduceAWindowEvenWithoutARing() { + // the interface default copies once; it must still describe the same positions + val turbo = sk.ainet.lang.tensor.storage.KvCacheStore.turboQuant( + numLayers = 1, numHeads = heads, headDim = headDim, maxSeqLen = 8, + ) + for (step in 0 until 3) turbo.appendToken(0, token(step), token(step, 0.5f)) + val w = turbo.keyWindow(0, 0, 3) + assertFalse(w.wrapped, "a compressed store hands back one run") + assertEquals(3, w.length) + assertEquals(heads, w.heads) + assertEquals(headDim, w.headDim) + } +}