feat(memory): KV cache preallocates its ring in Scope.Model (SKEEP-003 P2, S1.10)

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 <init> (Lsk/ainet/lang/tensor/storage/KvCacheConfig;)V
+	public fun <init> (Lsk/ainet/lang/tensor/storage/KvCacheConfig;Lsk/ainet/lang/memory/ModelScope;)V
+	public synthetic fun <init> (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;

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<FloatArray> = Array(numLayers) {
-        FloatArray(numHeads * maxSeqLen * headDim)
-    }
-    private val values: Array<FloatArray> = 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<FloatArray> = Array(numLayers) { allocateLayer("k", it) }
+    private val values: Array<FloatArray> = 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) {

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<TraceEvent.Allocation>()
+        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<TraceEvent.Allocation>().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<TraceEvent.Allocation>().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<TraceEvent.Free>().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)
+    }
+}