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) + } +}