From 0b6a23d2a4cbcab3412cf76ae3fa26e58b102658 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 21:36:00 +0200 Subject: [PATCH] feat(memory): KV cache preallocates its ring in Scope.Model (SKEEP-003 P2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F2. The cache already preallocated its whole ring; what it lacked was an owner. DefaultKvCacheStore now takes an optional ModelScope and allocates the per-layer K/V backing through it: - tracked and traced — one Allocation event per layer and side, in MODEL scope, carrying the cache's TensorId (kv.layers[N].k / .v) instead of an anonymous buffer, so the ring shows up in plan-vs-actual (#1074) as model-scope bytes and never as forward-scope churn; - released deterministically when the model closes (Free events summing to the ring size), instead of waiting for the GC; - preallocatedBytes exposes what the ring costs (layers × 2 × heads × maxSeqLen × headDim × 4) so a caller can compare it with the plan's KV line. Without a scope the store behaves exactly as before (plain arrays, GC lifetime), so nothing existing changes. The TurboQuant store keeps its block-encoded path; giving it the same owner is a follow-up. KvCacheModelScopeTest: the no-scope path is unchanged; with a scope the ring is allocated once per layer/side in MODEL scope with the right ids and total, appending tokens allocates nothing more, closing the model frees exactly the ring, and ActualMemory sees it as model-scope peak with zero forward bytes. 177/177 storage tests. Closes #1031 Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-lang-core.api | 3 + .../tensor/storage/DefaultKvCacheStore.kt | 38 +++++++-- .../tensor/storage/KvCacheModelScopeTest.kt | 82 +++++++++++++++++++ 3 files changed, 115 insertions(+), 8 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheModelScopeTest.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 c75821fad..a7c35d9ae 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 @@ -6306,6 +6306,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 appendToken (I[F[F)V public fun clear ()V public fun evict (I)V @@ -6316,6 +6318,7 @@ public final class sk/ainet/lang/tensor/storage/DefaultKvCacheStore : sk/ainet/l public fun getNumHeads ()I public fun getNumLayers ()I public fun getPlacement ()Lsk/ainet/lang/tensor/storage/Placement; + public final fun getPreallocatedBytes ()J public fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun memoryReport ()Lsk/ainet/lang/tensor/storage/KvCacheMemoryReport; public fun readKeyStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage; 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 8c145bf45..7f71b2203 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 @@ -17,8 +17,17 @@ import sk.ainet.lang.types.FP32 * * Append writes to position [currentSeqLen]; read returns a contiguous slice. */ -public class DefaultKvCacheStore( - private val config: KvCacheConfig +/** + * @param scope when given, the per-layer K/V backing is allocated in that [sk.ainet.lang.memory.ModelScope] + * instead of as plain GC-managed arrays (SKEEP-003 §4.5, PRD M1-F2): the allocation is tracked, + * traced (`TraceEvent.Allocation` with the cache's `TensorId`) and released deterministically when + * the model closes, which is what lets the memory plan be checked against reality (#1074). The + * default keeps today's behaviour exactly. + */ +@OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) +public class DefaultKvCacheStore @kotlin.jvm.JvmOverloads constructor( + private val config: KvCacheConfig, + private val scope: sk.ainet.lang.memory.ModelScope? = null, ) : KvCacheStore { override val numLayers: Int get() = config.numLayers @@ -33,14 +42,27 @@ public class DefaultKvCacheStore( override val currentSeqLen: Int get() = _currentSeqLen // Per-layer storage: keys[layer] and values[layer] - // Each is [numHeads, maxSeqLen, headDim] laid out as contiguous float array - private val keys: Array = Array(numLayers) { - FloatArray(numHeads * maxSeqLen * headDim) - } - private val values: Array = Array(numLayers) { - FloatArray(numHeads * maxSeqLen * headDim) + // Each is [numHeads, maxSeqLen, headDim] laid out as contiguous float array. + // With a ModelScope the arrays come from scope-owned storage (tracked, traced, freed with the + // model); without one they are plain arrays, exactly as before. + private val elementsPerLayer: Int = numHeads * maxSeqLen * headDim + + private fun allocateLayer(kind: String, layer: Int): FloatArray { + val s = scope ?: return FloatArray(elementsPerLayer) + val storage = s.allocateFloats( + elementsPerLayer, + sk.ainet.lang.tensor.TensorId(listOf("kv", "layers[$layer]"), kind), + ) + val array = storage.floats ?: FloatArray(elementsPerLayer) + return if (storage.arrayOffset == 0 && array.size == elementsPerLayer) array else FloatArray(elementsPerLayer) } + private val keys: Array = Array(numLayers) { allocateLayer("k", it) } + private val values: Array = Array(numLayers) { allocateLayer("v", it) } + + /** Bytes of K/V backing this store preallocated (`layers × 2 × heads × maxSeqLen × headDim × 4`). */ + public val preallocatedBytes: Long get() = numLayers.toLong() * 2 * elementsPerLayer * 4 + override fun appendToken(layer: Int, key: FloatArray, value: FloatArray) { requireLayerIndex(layer) check(_currentSeqLen < maxSeqLen) { diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheModelScopeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheModelScopeTest.kt new file mode 100644 index 000000000..a6f011997 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheModelScopeTest.kt @@ -0,0 +1,82 @@ +package sk.ainet.lang.tensor.storage + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.ModelScope +import sk.ainet.lang.memory.ScopeKind +import sk.ainet.lang.memory.plan.ActualMemory +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * SKEEP-003 §4.5 / PRD M1-F2: the KV cache preallocates its whole ring in `Scope.Model`, so the + * bytes are tracked, traced and released with the model — and the memory plan's KV line can be + * checked against what the store actually took (#1074). + */ +@OptIn(ExperimentalMemoryApi::class) +class KvCacheModelScopeTest { + + private val config = KvCacheConfig(numLayers = 4, numHeads = 2, headDim = 8, maxSeqLen = 16) + + private fun expectedBytes(): Long = 4L * 2 * (2 * 16 * 8) * 4 // layers × (K+V) × heads·seq·dim × 4 B + + @Test + fun withoutAScopeTheStoreBehavesExactlyAsBefore() { + val store = DefaultKvCacheStore(config) + assertEquals(expectedBytes(), store.preallocatedBytes) + store.appendToken(0, FloatArray(16) { 1f }, FloatArray(16) { 2f }) + for (l in 1 until 4) store.appendToken(l, FloatArray(16) { 1f }, FloatArray(16) { 2f }) + assertEquals(1, store.currentSeqLen) + assertEquals(1f, store.readKeys(0, 0, 1)[0]) + } + + @Test + fun withAModelScopeTheRingIsAllocatedTracedAndFreedWithTheModel() { + val sink = RecordingTraceSink() + val model = ModelScope(sink, "llama") + val store = DefaultKvCacheStore(config, model) + + // one allocation per layer per side, all in MODEL scope, summing to the ring size + val allocations = sink.eventsOf() + assertEquals(8, allocations.size, "4 layers × K and V") + assertTrue(allocations.all { it.scope == ScopeKind.MODEL }) + assertEquals(expectedBytes(), allocations.sumOf { it.bytes }) + assertEquals(expectedBytes(), store.preallocatedBytes) + assertEquals(expectedBytes(), model.liveBytes) + // ids name the cache, not an anonymous buffer + assertTrue(allocations.any { it.origin?.canonical == "kv.layers[0].k" }, allocations.map { it.origin?.canonical }.toString()) + assertTrue(allocations.any { it.origin?.canonical == "kv.layers[3].v" }) + + // the cache still works + for (l in 0 until 4) store.appendToken(l, FloatArray(16) { (l + 1).toFloat() }, FloatArray(16) { -1f }) + assertEquals(1, store.currentSeqLen) + assertEquals(3f, store.readKeys(2, 0, 1)[0]) + + // steady state: appending tokens allocates nothing more + val before = sink.eventsOf().size + for (step in 1 until 5) for (l in 0 until 4) store.appendToken(l, FloatArray(16) { 1f }, FloatArray(16) { 1f }) + assertEquals(before, sink.eventsOf().size, "appending must not allocate") + assertEquals(5, store.currentSeqLen) + + // closing the model frees the ring + model.close() + assertEquals(0L, model.liveBytes) + assertEquals(expectedBytes(), sink.eventsOf().sumOf { it.bytes }) + } + + @Test + fun theRingShowsUpInPlanVsActualAsModelScopeBytes() { + val sink = RecordingTraceSink() + val model = ModelScope(sink) + DefaultKvCacheStore(config, model) + val actual = ActualMemory.from(sink) + assertEquals(expectedBytes(), actual.peakModelBytes) + assertEquals(0L, actual.peakForwardBytes, "a KV ring is model-lifetime, never forward") + assertEquals(8, actual.allocationsByScope[ScopeKind.MODEL]) + model.close() + assertFalse(model.liveBytes > 0) + } +}