test(m1): synthetic decode harness asserting M1's memory acceptance criteria (SKEEP-003 M1, S1.11a)

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/DecodeAcceptanceTest.kt

@@ -0,0 +1,92 @@
+package sk.ainet.exec.harness
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.ScopeKind
+import sk.ainet.lang.memory.plan.ActualMemory
+import sk.ainet.lang.memory.plan.PlanVsActual
+import sk.ainet.lang.memory.trace.PerfettoTraceExporter
+import sk.ainet.lang.memory.trace.TraceEvent
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * Milestone M1's acceptance criteria, asserted on every commit against the synthetic decode harness
+ * (#1032 option (c)): the memory behaviour is checked here, where the memory code lives; the real
+ * model's tok/s and TTFT belong to `skainet-decode` in SKaiNET-transformers.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+class DecodeAcceptanceTest {
+
+    // Enough steps to prove the staircase is flat, few enough for Karma's 2 s per-test budget in a
+    // browser (the reference kernel decodes every element, so wasm is the slowest target here).
+    private val steps = 12
+
+    @Test
+    fun m1a1_memoryIsFlatAcrossDecodeSteps() {
+        val h = DecodeHarness()
+        h.decode(steps)
+        // live bytes after the last step must equal live bytes after warm-up: the forward scope is
+        // recycled, the model scope holds weights + KV and nothing else grows.
+        val live = h.liveBytes()
+        val model = live[ScopeKind.MODEL] ?: 0L
+        val forward = live[ScopeKind.FORWARD] ?: 0L
+        assertEquals(0L, forward, "the forward scope must be empty between steps")
+        assertTrue(model > 0, "weights and the KV ring stay resident")
+
+        // and the per-step resets all report the same live-bytes-before, i.e. a flat staircase
+        val resets = h.sink.eventsOf<TraceEvent.ScopeReset>()
+        assertEquals(steps, resets.size)
+        val warmed = resets.drop(3).map { it.liveBytesBefore }.toSet()
+        assertEquals(1, warmed.size, "per-step forward use must be identical after warm-up, saw $warmed")
+        assertTrue(resets.all { it.liveBytesAfter == 0L })
+        h.close()
+    }
+
+    @Test
+    fun m1a3_noForwardScopeAllocationsPerStepAfterWarmUp() {
+        val h = DecodeHarness()
+        h.decode(12)
+        // the slab is allocated once, before the first step; steps 5..20 must add nothing
+        assertEquals(0, h.allocationsBetweenSteps(ScopeKind.FORWARD, fromStep = 4, toStep = 12), "steady-state decode must not allocate")
+        assertEquals(0, h.allocationsBetweenSteps(ScopeKind.MODEL, fromStep = 1, toStep = 12), "weights and KV are allocated before decoding")
+        h.close()
+    }
+
+    @Test
+    fun m1a8_planMatchesWhatTheRunAllocated() {
+        val h = DecodeHarness()
+        h.decode(12)
+        val cmp = PlanVsActual.of(h.plan(), h.sink)
+        assertTrue(cmp.withinTolerance, "plan drifted from the run:\n" + cmp.render())
+        val actual = ActualMemory.from(h.sink)
+        assertTrue(actual.peakModelBytes > 0)
+        assertEquals(0L, actual.adapterBytes, "a well-formed decode step needs no adapters")
+        h.close()
+    }
+
+    @Test
+    fun m1a7_theTraceHasOneTrackPerScopeAndAFlatLiveBytesCounter() {
+        val h = DecodeHarness()
+        h.decode(10)
+        val json = PerfettoTraceExporter.export(h.sink, processName = "skainet-decode-harness")
+        assertTrue(json.contains("model scope") && json.contains("forward scope"), "one track per scope")
+        assertTrue(json.contains("\"ph\":\"X\",\"name\":\"matmul\""), "kernel spans")
+        assertTrue(json.contains("model.layers[0].attn.q_proj.weight"), "spans labelled by TensorId")
+        assertTrue(json.contains("\"ph\":\"C\",\"name\":\"live bytes\""), "a live-bytes counter track")
+        assertTrue(json.contains("\"forward\":0"), "the counter returns to zero at every reset")
+        h.close()
+    }
+
+    @Test
+    fun theHarnessActuallyExercisesTheMemoryModel() {
+        val h = DecodeHarness()
+        h.decode(3)
+        val kernels = h.sink.eventsOf<TraceEvent.KernelRun>()
+        assertTrue(kernels.isNotEmpty(), "matmuls must go through KernelDispatch")
+        assertTrue(kernels.all { it.op == "matmul" })
+        assertTrue(h.sink.eventsOf<TraceEvent.Allocation>().any { it.scope == ScopeKind.MODEL })
+        assertTrue(h.sink.eventsOf<TraceEvent.PhaseBegin>().count { it.phase == "decode" } == 3)
+        h.close()
+    }
+}

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/DecodeHarness.kt

@@ -0,0 +1,154 @@
+package sk.ainet.exec.harness
+
+import sk.ainet.backend.api.kernel.KernelDispatch
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.ForwardScope
+import sk.ainet.lang.memory.ModelScope
+import sk.ainet.lang.memory.PackedBlockDecoder
+import sk.ainet.lang.memory.ScopeKind
+import sk.ainet.lang.memory.Storage
+import sk.ainet.lang.memory.TensorView
+import sk.ainet.lang.memory.plan.KvCacheMode
+import sk.ainet.lang.memory.plan.MemoryPlans
+import sk.ainet.lang.memory.plan.ModelGeometry
+import sk.ainet.lang.memory.plan.PlanInput
+import sk.ainet.lang.memory.plan.PlanTensor
+import sk.ainet.lang.memory.trace.RecordingTraceSink
+import sk.ainet.lang.memory.trace.TraceEvent
+import sk.ainet.lang.memory.trace.phase
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.TensorId
+import sk.ainet.lang.tensor.data.Q8_0BlockTensorData
+import sk.ainet.lang.tensor.storage.DefaultKvCacheStore
+import sk.ainet.lang.tensor.storage.KvCacheConfig
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+
+/**
+ * A synthetic decode loop over the *real* memory machinery — SKEEP-003 M1's acceptance harness
+ * (#1032, option (c)).
+ *
+ * It is deliberately not a model: no tokenizer, no sampling, no checkpoint. It is a Llama-shaped
+ * stack of packed matmuls whose weights live in a [ModelScope] (mapped-equivalent, packed Q8_0),
+ * whose activations come from a recycled [ForwardScope], whose KV ring is preallocated in the model
+ * scope, and whose dispatch goes through [KernelDispatch] — so the criteria that are about *memory
+ * behaviour* can be asserted on every commit in the repository where that behaviour lives:
+ *
+ * - **M1-A1** flat memory: live bytes at the last step equal live bytes after warm-up;
+ * - **M1-A3** zero forward-scope allocations per step after warm-up;
+ * - **M1-A8** the plan matches what the run allocated;
+ * - **M1-A7** the trace has one track per scope and a live-bytes counter that returns to zero.
+ *
+ * The real end-to-end numbers (tok/s, TTFT, peak load RSS, effective bandwidth) belong to the
+ * `skainet-decode` sample in SKaiNET-transformers, which owns the model.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+public class DecodeHarness(
+    // Deliberately tiny: the reference kernel decodes every element, and these tests run in a
+    // browser under Karma's 2 s per-test budget as well as on the JVM. The assertions are about
+    // *memory behaviour*, which does not need big shapes — the real numbers come from
+    // skainet-decode in SKaiNET-transformers.
+    public val layers: Int = 2,
+    public val hidden: Int = 32,
+    public val ffn: Int = 64,
+    public val heads: Int = 4,
+    public val kvHeads: Int = 2,
+    public val ctx: Int = 32,
+    public val vocab: Int = 64,
+) {
+    public val sink: RecordingTraceSink = RecordingTraceSink()
+    private val model = ModelScope(sink, "harness")
+    private val blockSize = 32
+
+    /** Per-layer weights, packed Q8_0, living in the model scope (the mapped-weight shape). */
+    private val weights: List<TensorView> = buildList {
+        for (l in 0 until layers) {
+            add(packedWeight(hidden, hidden, TensorId(listOf("model", "layers[$l]", "attn"), "q_proj.weight")))
+            add(packedWeight(ffn, hidden, TensorId(listOf("model", "layers[$l]", "mlp"), "up_proj.weight")))
+        }
+    }
+
+    private val kv = DefaultKvCacheStore(
+        KvCacheConfig(numLayers = layers, numHeads = kvHeads, headDim = hidden / heads, maxSeqLen = ctx),
+        model,
+    )
+
+    /** The forward slab: activations of one step, sized from the plan. */
+    private val forward = ForwardScope(slabFloats = forwardFloats(), sink = sink, name = "decode")
+
+    private fun forwardFloats(): Int = (4 * hidden + 3 * ffn + heads * ctx) + vocab
+
+    private fun packedWeight(rows: Int, cols: Int, id: TensorId): TensorView {
+        require(cols % blockSize == 0)
+        val blocks = rows * (cols / blockSize)
+        val bytes = ByteArray(blocks * 34)
+        var seed = id.canonical.hashCode()
+        for (b in 0 until blocks) {
+            val off = b * 34
+            // a sane FP16 scale, then deterministic codes
+            bytes[off] = 0x00; bytes[off + 1] = 0x38          // half(0.5)
+            for (i in 0 until 32) { seed = seed * 1103515245 + 12345; bytes[off + 2 + i] = (seed ushr 16).toByte() }
+        }
+        val storage = model.adopt(Storage.Heap.wrap(bytes, mutable = false, origin = id, sink = sink))
+        val data = Q8_0BlockTensorData(Shape(rows, cols), bytes)
+        return TensorView.packed(storage, Shape(rows, cols), TensorEncoding.Q8_0, PackedBlockDecoder(data), id = id)
+    }
+
+    /** The memory plan this harness's shapes predict. */
+    public fun plan(): sk.ainet.lang.memory.plan.MemoryPlan {
+        val f = Format(FP32, TensorEncoding.Q8_0)
+        val tensors = weights.map { w ->
+            PlanTensor(w.id!!.canonical, w.id, f, w.elementCount, f.physicalBytes(w.elementCount)!!)
+        }
+        val geometry = ModelGeometry(layers, heads, kvHeads, hidden / heads, hidden / heads, hidden, ffn, vocab)
+        return MemoryPlans.plan(PlanInput("harness", "llama", tensors, geometry, ctx, prefillChunk = 1, kvMode = KvCacheMode.FP32))
+    }
+
+    /** Run [steps] decode steps; each allocates activations, runs the stack and resets the scope. */
+    public fun decode(steps: Int) {
+        val x = FloatArray(hidden) { (it % 7) * 0.125f }
+        for (step in 1..steps) {
+            sink.phase("decode", step) {
+                val act = forward.allocateFloats(hidden, TensorId(listOf("model"), "hidden", "step=$step"))
+                x.copyInto(act.floats!!, act.arrayOffset)
+                val actView = TensorView.dense(act, Shape(1, hidden), FP32, TensorId(listOf("model"), "hidden", "step=$step"))
+                for (w in weights) {
+                    if (w.shape[1] != hidden) continue
+                    val out = forward.allocateFloats(w.shape[0], TensorId(listOf("model"), "proj", "step=$step"))
+                    val outView = TensorView.dense(out, Shape(1, w.shape[0]), FP32)
+                    KernelDispatch.matmul(actView, w, outView, forward, sink)
+                }
+                // one token into the KV ring (all layers), as a decode step does
+                val k = FloatArray(kvHeads * (hidden / heads)) { 0.5f }
+                if (kv.currentSeqLen < ctx) for (l in 0 until layers) kv.appendToken(l, k, k)
+                forward.reset()
+            }
+        }
+    }
+
+    /** Live bytes per scope as the event stream saw them after the last step. */
+    public fun liveBytes(): Map<ScopeKind, Long> {
+        val live = HashMap<ScopeKind, Long>()
+        for (e in sink.events()) when (e) {
+            is TraceEvent.Allocation -> live[e.scope] = (live[e.scope] ?: 0L) + e.bytes
+            is TraceEvent.Free -> live[e.scope] = ((live[e.scope] ?: 0L) - e.bytes).coerceAtLeast(0L)
+            is TraceEvent.ScopeReset -> live[e.scope] = e.liveBytesAfter
+            else -> Unit
+        }
+        return live
+    }
+
+    /** Allocation events recorded in [scope] between two step numbers (inclusive of the phases). */
+    public fun allocationsBetweenSteps(scope: ScopeKind, fromStep: Int, toStep: Int): Int {
+        var step = 0
+        var count = 0
+        for (e in sink.events()) {
+            if (e is TraceEvent.PhaseBegin && e.phase == "decode") step = e.step ?: step
+            if (e is TraceEvent.Allocation && e.scope == scope && step in fromStep..toStep) count++
+        }
+        return count
+    }
+
+    public fun close() { forward.close(); model.close() }
+}

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -1212,6 +1212,7 @@ public final class sk/ainet/lang/memory/plan/Budget$Companion {
 
 public final class sk/ainet/lang/memory/plan/KvCacheMode : java/lang/Enum {
 	public static final field BF16 Lsk/ainet/lang/memory/plan/KvCacheMode;
+	public static final field FP32 Lsk/ainet/lang/memory/plan/KvCacheMode;
 	public static final field TURBOQUANT_4 Lsk/ainet/lang/memory/plan/KvCacheMode;
 	public final fun bytes (J)J
 	public static fun getEntries ()Lkotlin/enums/EnumEntries;

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Scope.kt

@@ -67,14 +67,33 @@ public class ModelScope(override val sink: TraceSink = NoopTraceSink, public val
     /** Map a file region (packed weights) into this scope; unmapped when the scope closes. */
     public fun mapFile(path: String, fileOffset: Long, length: Long, origin: TensorId? = null): Storage = track(PlatformStorage.mapFile(path, fileOffset, length, ScopeKind.MODEL, origin, sink))
 
-    /** Adopt a storage allocated elsewhere (e.g. by a loader) so the scope closes it. */
-    public fun adopt(storage: Storage): Storage = track(storage)
+    /**
+     * Adopt a storage allocated elsewhere (a loader's mapped weight, a borrowed array) so the scope
+     * closes it — and **count it**: this is the point where the model takes responsibility for those
+     * bytes. Mapped and borrowed weights emit no allocation event of their own (borrowing is not an
+     * allocation), yet decision #11 counts weights as resident because decode touches every weight
+     * every token, so without this the memory plan would be checked against a run that appears to
+     * hold nothing (#1074).
+     */
+    public fun adopt(storage: Storage): Storage {
+        val s = track(storage)
+        if (sink.isEnabled) sink.emit(TraceEvent.Allocation(s.id.value, ScopeKind.MODEL, s.sizeBytes, s.debugOrigin, site = "adopted"))
+        return s
+    }
 
     /** Close every owned storage (weights unmapped, off-heap freed) — exactly once. */
     override fun close() {
         if (closed) return
         closed = true
-        for (s in owned.asReversed()) s.close()
+        for (s in owned.asReversed()) {
+            val wasBorrowed = s.owner is Owner.Borrowed
+            val bytes = s.sizeBytes
+            val id = s.id
+            s.close()
+            // A borrowed storage does not emit Free on its own (we never freed anything), but the
+            // scope did stop counting it — keep the ledger balanced for plan-vs-actual.
+            if (wasBorrowed && sink.isEnabled) sink.emit(TraceEvent.Free(id.value, ScopeKind.MODEL, bytes))
+        }
         owned.clear()
     }
 }

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt

@@ -47,13 +47,21 @@ public data class ModelGeometry(
 /** How the KV cache is stored. */
 @ExperimentalMemoryApi
 public enum class KvCacheMode(public val label: String) {
-    /** 2 bytes per element (bf16/f16). */
+    /**
+     * 4 bytes per element — what `DefaultKvCacheStore` actually stores today (dense FloatArray
+     * rings). Planning a dense store as [BF16] understates it by 2×, which is the kind of drift
+     * the plan-vs-actual check exists to catch (#1074).
+     */
+    FP32("fp32"),
+
+    /** 2 bytes per element (bf16/f16) — a narrow-float KV store. */
     BF16("bf16"),
     /** TurboQuant 4-bit polar codes, block 128 (decision #11 default when the plan is tight). */
     TURBOQUANT_4("TurboQuant 4-bit");
 
     /** Bytes for [elements] cache elements under this mode. */
     public fun bytes(elements: Long): Long = when (this) {
+        FP32 -> 4L * elements
         BF16 -> 2L * elements
         TURBOQUANT_4 -> TensorEncoding.TurboQuantPolar(bitsPerElement = 4, blockSize = 128).physicalBytes(elements) ?: (elements / 2)
     }
@@ -164,7 +172,7 @@ public data class MemoryPlan(
         for (l in lines) {
             append("  "); append(l.section.padEnd(10)); append(l.detail.padEnd(26)); append(MemoryPlans.formatBytes(l.bytes).padStart(10))
             if (l.resident) append("   resident")
-            if (l.section == "kv cache") append("   (").append(MemoryPlans.formatBytes(kvBytesAlternate)).append(" with ").append(if (input.kvMode == KvCacheMode.BF16) KvCacheMode.TURBOQUANT_4.label else KvCacheMode.BF16.label).append(')')
+            if (l.section == "kv cache") append("   (").append(MemoryPlans.formatBytes(kvBytesAlternate)).append(" with ").append(if (input.kvMode == KvCacheMode.TURBOQUANT_4) KvCacheMode.BF16.label else KvCacheMode.TURBOQUANT_4.label).append(')')
             append('\n')
         }
         append("  "); append("total".padEnd(36)); append(MemoryPlans.formatBytes(totalBytes).padStart(10))
@@ -205,7 +213,7 @@ public object MemoryPlans {
         val g = input.geometry
         val kvElements = if (g != null) kvElements(g, input.ctx) else 0L
         val kv = input.kvMode.bytes(kvElements)
-        val kvAlt = (if (input.kvMode == KvCacheMode.BF16) KvCacheMode.TURBOQUANT_4 else KvCacheMode.BF16).bytes(kvElements)
+        val kvAlt = (if (input.kvMode == KvCacheMode.TURBOQUANT_4) KvCacheMode.BF16 else KvCacheMode.TURBOQUANT_4).bytes(kvElements)
         val forward = if (g != null) forwardBytes(g, input.ctx, input.prefillChunk) else 0L
         return MemoryPlan(input, weights, kv, kvAlt, forward, HEAP_HEADROOM_BYTES, budget)
     }