Perf/llama packed load memory 0.34

llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/GenerateExtensions.kt

@@ -14,9 +14,9 @@ public fun <T : DType> sampleFromTensor(
     logits: Tensor<T, Float>,
     temperature: Float,
     random: Random = Random.Default
-): Int {
+): Int = sk.ainet.lang.nn.transformer.PhaseProfile.time("sample") {
     val buf = extractFloatBuffer(logits).copyOf()
-    return sampleFromLogits(buf, temperature, random)
+    sampleFromLogits(buf, temperature, random)
 }
 
 /**

llm-inference/llama/src/androidMain/kotlin/sk/ainet/models/llama/GcHint.android.kt

@@ -0,0 +1,4 @@
+package sk.ainet.models.llama
+
+// Android (ART) GC is automatic — no-op (distinct source set from jvmMain).
+internal actual fun gcCollectHint() {}

llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt

@@ -331,16 +331,33 @@ public class DecoderGgufWeightLoader private constructor(
                 // the physical tensor is TOKEN_EMBEDDINGS — both must have [vocab, dim] shape.
                 validateStreamingTensorShape(name, st, metadata)
                 val tensor: Tensor<T, V> = streamingTensorToTensor(ctx, dtype, reader, st)
-                onTensorLoaded(name, tensor)
                 val isRawQuant = when (quantPolicy) {
                     QuantPolicy.RAW_BYTES -> st.tensorType != GGMLQuantizationType.F32
                     QuantPolicy.NATIVE_OPTIMIZED -> st.tensorType != GGMLQuantizationType.F32
                         && st.tensorType != GGMLQuantizationType.F16
                         && st.tensorType != GGMLQuantizationType.BF16
                     QuantPolicy.DEQUANTIZE_TO_FP32 -> false
                 }
-                if (isRawQuant) {
-                    quantCallback?.invoke(name, st.tensorType)
+                if (quantPolicy == QuantPolicy.NATIVE_OPTIMIZED && isRawQuant) {
+                    // Fuse load+pack: convert this tensor to its packed/FP32 form NOW so the raw
+                    // bytes are never accumulated in the map. Holding the whole raw model AND the
+                    // whole packed model at once is the ~3.2 GB load peak that OOM-kills the 2 GB
+                    // board (see PERF-LOGBOOK p1-result). The packed data encodes its own quant
+                    // type, so we DON'T register it in quantTypes — the post-load
+                    // convertLlamaWeightsPacked then leaves it untouched (idempotent).
+                    @Suppress("UNCHECKED_CAST")
+                    val packed = packLlamaTensor(
+                        name, tensor as Tensor<DType, Any>, st.tensorType, metadata, ctx,
+                    ) as Tensor<T, V>
+                    onTensorLoaded(name, packed)
+                    // Free this tensor's transient copies (raw Int8 + extractRawBytes + relayout)
+                    // before the next tensor; on Kotlin/Native they otherwise pile up uncollected.
+                    gcCollectHint()
+                } else {
+                    onTensorLoaded(name, tensor)
+                    if (isRawQuant) {
+                        quantCallback?.invoke(name, st.tensorType)
+                    }
                 }
             }
 

llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/GcHint.kt

@@ -0,0 +1,15 @@
+package sk.ainet.models.llama
+
+/**
+ * Hint the platform to reclaim unreachable memory *now*.
+ *
+ * On **Kotlin/Native** the GC is allocation-triggered and lags during a tight load/convert loop, so
+ * the per-tensor transient copies (`pread` ByteArray → `copyOf` → `extractRawBytes` → relayout
+ * ByteArray) pile up uncollected and inflate peak RSS — which OOM-kills the 2 GB board even though the
+ * resident model is only ~0.9 GB. Calling this at per-tensor boundaries (after the source tensor has
+ * been dropped) bounds the high-water to roughly one tensor in flight. Mirrors the established board
+ * pattern in `GemmaDecoder` (`kotlin.native.runtime.GC.collect()`).
+ *
+ * **No-op on the JVM** — the JVM GC reclaims under allocation pressure on its own.
+ */
+internal expect fun gcCollectHint()

llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaPackedWeights.kt

@@ -32,37 +32,65 @@ public fun convertLlamaWeightsPacked(
     val quantTypes = typed.quantTypes
     if (quantTypes.isEmpty()) return weights
 
+    // Memory: drain the SOURCE map as we convert (P1). The streaming loader builds
+    // `tensors` as a LinkedHashMap, so we can `remove` each raw tensor right after
+    // packing it — its bytes become GC-eligible immediately instead of the whole raw
+    // model staying resident alongside the whole packed model. That load-time doubling
+    // (~3.2 GB peak for TinyLlama Q4_K) is what OOM-kills the 2 GB board; draining keeps
+    // peak ≈ packed-so-far + the one tensor in flight. Falls back to non-destructive if
+    // the map is somehow immutable (correctness unchanged, just no memory win).
+    @Suppress("UNCHECKED_CAST")
+    val src = typed.tensors as? MutableMap<String, Tensor<DType, Any>>
+    val names = typed.tensors.keys.toList()
+
     val newTensors = linkedMapOf<String, Tensor<DType, Any>>()
-    for ((name, tensor) in typed.tensors) {
-        val qt = quantTypes[name]
-        newTensors[name] = when {
-            qt == null -> tensor // not quantized (norms, f32)
-            else -> {
-                val shape = logicalShapeFor(name, typed.metadata)
-                if (shape == null) {
-                    tensor // unknown 2-D layout — leave as-is
-                } else {
-                    val bytes = extractRawBytes(tensor.data)
-                    // token_embd is gathered (row lookup) → must be FP32. Other matrices (incl.
-                    // output/lm_head) stay packed and run the in-kernel matmul.
-                    val isEmbed = name == LlamaTensorNames.TOKEN_EMBEDDINGS
-                    val packed = if (!isEmbed) packLlamaKQuant<FP32>(bytes, qt, shape) else null
-                    when {
-                        packed != null -> {
-                            @Suppress("UNCHECKED_CAST")
-                            ctx.fromData(packed as TensorData<FP32, Float>, FP32::class) as Tensor<DType, Any>
-                        }
-                        isEmbed -> dequantNoTranspose(bytes, qt, shape, ctx)
-                        else -> dequantTransposed(bytes, qt, shape, ctx)
-                    }
-                }
-            }
-        }
+    for (name in names) {
+        val tensor = src?.remove(name) ?: typed.tensors.getValue(name)
+        // qt == null (norms/f32) or already-packed (streaming fused path leaves them out of
+        // quantTypes) → packLlamaTensor returns the tensor unchanged (idempotent).
+        newTensors[name] = packLlamaTensor(name, tensor, quantTypes[name], typed.metadata, ctx)
+        // Reclaim this tensor's transient copies (raw source + extractRawBytes + relayout input)
+        // before moving on — on Kotlin/Native they otherwise accumulate uncollected and blow up
+        // peak RSS. No-op on the JVM. Paired with the source-map drain above.
+        gcCollectHint()
     }
     @Suppress("UNCHECKED_CAST")
     return DecoderGgufWeights(typed.metadata, newTensors, typed.quantTypes) as DecoderGgufWeights<*, *>
 }
 
+/**
+ * Convert ONE raw NATIVE_OPTIMIZED tensor to its DSL-consumed form: Q4_K/Q5_K/Q6_K/Q8_0 matmul
+ * weights → heap-packed `Q*BlockTensorData`; `token_embd` → FP32 `[vocab, embed]` (gathered);
+ * other quantized-without-a-packed-kernel → FP32 transposed `[out, in]`.
+ *
+ * Shared by [convertLlamaWeightsPacked] (post-load pass, non-streaming) and the streaming loader's
+ * fused load+pack path. Idempotent: `qt == null` or an unknown 2-D layout returns [tensor] as-is,
+ * so packing the same map twice (or a map of already-packed tensors) is a no-op.
+ */
+internal fun packLlamaTensor(
+    name: String,
+    tensor: Tensor<DType, Any>,
+    qt: GGMLQuantizationType?,
+    metadata: LlamaModelMetadata,
+    ctx: ExecutionContext,
+): Tensor<DType, Any> {
+    if (qt == null) return tensor // not quantized (norms, f32), or already packed
+    val shape = logicalShapeFor(name, metadata) ?: return tensor // unknown 2-D layout — leave as-is
+    val bytes = extractRawBytes(tensor.data)
+    // token_embd is gathered (row lookup) → must be FP32. Other matrices (incl. output/lm_head)
+    // stay packed and run the in-kernel matmul.
+    val isEmbed = name == LlamaTensorNames.TOKEN_EMBEDDINGS
+    val packed = if (!isEmbed) packLlamaKQuant<FP32>(bytes, qt, shape) else null
+    return when {
+        packed != null -> {
+            @Suppress("UNCHECKED_CAST")
+            ctx.fromData(packed as TensorData<FP32, Float>, FP32::class) as Tensor<DType, Any>
+        }
+        isEmbed -> dequantNoTranspose(bytes, qt, shape, ctx)
+        else -> dequantTransposed(bytes, qt, shape, ctx)
+    }
+}
+
 /** Dequant to FP32 in natural `[rows, cols]` order (embeddings — gathered, not matmul'd). */
 @Suppress("UNCHECKED_CAST")
 private fun dequantNoTranspose(

llm-inference/llama/src/jsMain/kotlin/sk/ainet/models/llama/GcHint.js.kt

@@ -0,0 +1,4 @@
+package sk.ainet.models.llama
+
+// JS / Wasm GC is automatic and these targets don't run the board load path — no-op.
+internal actual fun gcCollectHint() {}

llm-inference/llama/src/jvmMain/kotlin/sk/ainet/models/llama/GcHint.jvm.kt

@@ -0,0 +1,4 @@
+package sk.ainet.models.llama
+
+// JVM GC reclaims under allocation pressure on its own — no explicit collect needed.
+internal actual fun gcCollectHint() {}

llm-inference/llama/src/nativeMain/kotlin/sk/ainet/models/llama/GcHint.native.kt

@@ -0,0 +1,6 @@
+package sk.ainet.models.llama
+
+@OptIn(kotlin.native.runtime.NativeRuntimeApi::class)
+internal actual fun gcCollectHint() {
+    kotlin.native.runtime.GC.collect()
+}

llm-inference/llama/src/wasmJsMain/kotlin/sk/ainet/models/llama/GcHint.wasmJs.kt

@@ -0,0 +1,4 @@
+package sk.ainet.models.llama
+
+// JS / Wasm GC is automatic and these targets don't run the board load path — no-op.
+internal actual fun gcCollectHint() {}

llm-inference/llama/src/wasmWasiMain/kotlin/sk/ainet/models/llama/GcHint.wasmWasi.kt

@@ -0,0 +1,4 @@
+package sk.ainet.models.llama
+
+// JS / Wasm GC is automatic and these targets don't run the board load path — no-op.
+internal actual fun gcCollectHint() {}

transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/RMSNormalization.kt

@@ -87,25 +87,27 @@ public class RMSNormalization<T : DType, V>(
 
     override fun forward(input: Tensor<T, V>, ctx: ExecutionContext): Tensor<T, V> =
         sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) {
-            // Try fused path if the ops backend supports it. Skip for Gemma's
-            // unit-offset variant — backends don't know about (1 + weight).
-            val w = params[0].value
-            if (!unitOffset) {
-                val fusedOps = ctx.ops as? FusedRmsNormOps
-                if (fusedOps != null) {
-                    val result = fusedOps.fusedRmsNorm(input, w, eps.toFloat())
-                    if (result != null) return@withForwardHooks result
+            sk.ainet.lang.nn.transformer.PhaseProfile.time("rmsnorm") {
+                // Try fused path if the ops backend supports it. Skip for Gemma's
+                // unit-offset variant — backends don't know about (1 + weight).
+                val w = params[0].value
+                val fused = if (!unitOffset) {
+                    (ctx.ops as? FusedRmsNormOps)?.fusedRmsNorm(input, w, eps.toFloat())
+                } else null
+                if (fused != null) {
+                    fused
+                } else {
+                    // Fallback: decomposed path
+                    val squared = input * input
+                    val mean = squared.mean(dim = input.rank - 1)
+                    // Unsqueeze so broadcasting works for batched input (e.g. [B, dim] / [B, 1])
+                    val rmsRaw = (mean + eps).sqrt()
+                    val rms = if (rmsRaw.rank < input.rank) rmsRaw.unsqueeze(rmsRaw.rank) else rmsRaw
+                    val normalized = input / rms
+                    val gain = if (unitOffset) ctx.ops.addScalar(w, 1f) else w
+                    val gainBcast = if (gain.rank == 1) gain.reshape(Shape(1, gain.shape[0])) else gain
+                    normalized * gainBcast
                 }
             }
-            // Fallback: decomposed path
-            val squared = input * input
-            val mean = squared.mean(dim = input.rank - 1)
-            // Unsqueeze so broadcasting works for batched input (e.g. [B, dim] / [B, 1])
-            val rmsRaw = (mean + eps).sqrt()
-            val rms = if (rmsRaw.rank < input.rank) rmsRaw.unsqueeze(rmsRaw.rank) else rmsRaw
-            val normalized = input / rms
-            val gain = if (unitOffset) ctx.ops.addScalar(w, 1f) else w
-            val gainBcast = if (gain.rank == 1) gain.reshape(Shape(1, gain.shape[0])) else gain
-            normalized * gainBcast
         }
 }

transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt

@@ -241,9 +241,9 @@ public class MultiHeadAttention<T : DType, V>(
         // Project Q, K, V: input @ W^T (+ bias if enabled).
         // Q always comes from qInput; in self-attn kvInput === qInput,
         // in cross-attn kvInput is the encoder memory (different seqLen).
-        var q = linearProject(ops, qInput, wQ)
-        var k = linearProject(ops, kvInput, wK)
-        var v = linearProject(ops, kvInput, wV)
+        var q = PhaseProfile.time("attn.qkv_proj") { linearProject(ops, qInput, wQ) }
+        var k = PhaseProfile.time("attn.qkv_proj") { linearProject(ops, kvInput, wK) }
+        var v = PhaseProfile.time("attn.qkv_proj") { linearProject(ops, kvInput, wV) }
         if (mhaDump) {
             mhaDumpStat("[blk.0.mha post-Q-proj      ]", q)
             mhaDumpStat("[blk.0.mha post-K-proj      ]", k)
@@ -277,9 +277,9 @@ public class MultiHeadAttention<T : DType, V>(
         // can't reuse it here; we materialise the permute via a copy.
         val qSeqLen = if (qInput.rank >= 2) qInput.shape[qInput.rank - 2] else 1
         val kvSeqLen = if (kvInput.rank >= 2) kvInput.shape[kvInput.rank - 2] else 1
-        q = swapSeqHeadDims(ops.reshape(q, Shape(qSeqLen, nHeads, headDim)), ctx)
-        k = swapSeqHeadDims(ops.reshape(k, Shape(kvSeqLen, nKVHeads, headDim)), ctx)
-        var vReshaped = swapSeqHeadDims(ops.reshape(v, Shape(kvSeqLen, nKVHeads, headDim)), ctx)
+        q = PhaseProfile.time("attn.reshape") { swapSeqHeadDims(ops.reshape(q, Shape(qSeqLen, nHeads, headDim)), ctx) }
+        k = PhaseProfile.time("attn.reshape") { swapSeqHeadDims(ops.reshape(k, Shape(kvSeqLen, nKVHeads, headDim)), ctx) }
+        var vReshaped = PhaseProfile.time("attn.reshape") { swapSeqHeadDims(ops.reshape(v, Shape(kvSeqLen, nKVHeads, headDim)), ctx) }
 
         // Optional QK-Norm
         if (qNorm != null && kNorm != null) {
@@ -311,8 +311,8 @@ public class MultiHeadAttention<T : DType, V>(
         if (ropeModule != null && !isCrossAttention) {
             val position = kvCache?.position ?: 0
             if (mhaDump) println("[blk.0.mha pos=$position]")
-            q = ropeModule.forward(q, position, ctx)
-            k = ropeModule.forward(k, position, ctx)
+            q = PhaseProfile.time("attn.rope") { ropeModule.forward(q, position, ctx) }
+            k = PhaseProfile.time("attn.rope") { ropeModule.forward(k, position, ctx) }
             if (mhaDump) {
                 mhaDumpStat("[blk.0.mha post-RoPE-Q      ]", q)
                 mhaDumpStat("[blk.0.mha post-RoPE-K      ]", k)
@@ -324,7 +324,7 @@ public class MultiHeadAttention<T : DType, V>(
         // different ownership). Caching is the runtime's responsibility
         // for cross-attention and is rejected at the entry above.
         val (fullK, fullV) = if (kvCache != null && !isCrossAttention) {
-            kvCache!!.update(k, vReshaped, ctx)
+            PhaseProfile.time("attn.kvcache") { kvCache!!.update(k, vReshaped, ctx) }
         } else {
             k to vReshaped
         }
@@ -345,7 +345,7 @@ public class MultiHeadAttention<T : DType, V>(
         // max-stable softmax, same GQA head mapping head h → kv head h/nRep).
         if (qSeqLen == 1 && !isCrossAttention && slidingWindow == null) {
             val merged = fusedDecodeAttention(q, fullK, fullV, scale, ctx)
-            var output = linearProject(ops, merged, wO)
+            var output = PhaseProfile.time("attn.o_proj") { linearProject(ops, merged, wO) }
             if (bias) output = ops.add(output, params[oWIdx + 1].value)
             if (mhaDump) mhaDumpStat("[blk.0.mha post-fused-decode ]", output)
             return output
@@ -425,13 +425,17 @@ public class MultiHeadAttention<T : DType, V>(
         scale: Float,
         ctx: ExecutionContext,
     ): Tensor<T, V> {
-        val qBuf = q.data.copyToFloatArray()        // [nHeads * headDim]
-        val kBuf = fullK.data.copyToFloatArray()    // [nKVHeads * seqKV * headDim]
-        val vBuf = fullV.data.copyToFloatArray()    // [nKVHeads * seqKV * headDim]
+        // The three buffer copies grow with the KV length (K/V are the full
+        // cache) and repeat every token × layer — timed separately from the
+        // arithmetic so the profile can tell copy cost from compute cost.
+        val qBuf = PhaseProfile.time("attn.fused_copy") { q.data.copyToFloatArray() }        // [nHeads * headDim]
+        val kBuf = PhaseProfile.time("attn.fused_copy") { fullK.data.copyToFloatArray() }    // [nKVHeads * seqKV * headDim]
+        val vBuf = PhaseProfile.time("attn.fused_copy") { fullV.data.copyToFloatArray() }    // [nKVHeads * seqKV * headDim]
         val seqKV = fullK.shape[1]
         val nRep = nHeads / nKVHeads
         val out = FloatArray(nHeads * headDim)      // == qDim, row-major [h, d]
         val scores = FloatArray(seqKV)
+        PhaseProfile.time("attn.fused_compute") {
         for (h in 0 until nHeads) {
             val g = h / nRep                        // GQA: which KV head this query head reads
             val qOff = h * headDim
@@ -464,6 +468,7 @@ public class MultiHeadAttention<T : DType, V>(
                 out[oOff + d] = acc * inv
             }
         }
+        }
         @Suppress("UNCHECKED_CAST")
         return ctx.fromData(
             sk.ainet.lang.tensor.data.DenseFloatArrayTensorData<T>(Shape(1, qDim), out)

transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/PhaseProfile.kt

@@ -0,0 +1,54 @@
+package sk.ainet.lang.nn.transformer
+
+import kotlin.time.TimeSource
+
+/**
+ * Lightweight, always-on accumulating profiler for the eager decode phases.
+ * Diagnostic only — the sibling of the core backend's `KernelProfile` (which
+ * times the matmul dispatch paths); this one attributes the NON-matmul decode
+ * tail: attention plumbing, RMSNorm, RoPE, KV-cache, sampling, detok.
+ *
+ * Named-bucket design: `time("attn.rope") { ... }` accumulates wall time and a
+ * call count per bucket. The map lookup + lambda per call is negligible next to
+ * millisecond-scale phases (a few hundred calls per generated token). The
+ * decode path is single-threaded, so plain mutable state is safe.
+ *
+ * Read [report] after a run; [reset] between phases (e.g. after model load so
+ * the breakdown covers only prefill+decode).
+ */
+public object PhaseProfile {
+    private val clock = TimeSource.Monotonic
+    private val nanos = LinkedHashMap<String, Long>()
+    private val calls = LinkedHashMap<String, Long>()
+
+    public fun <R> time(phase: String, body: () -> R): R {
+        val mark = clock.markNow()
+        val r = body()
+        val ns = mark.elapsedNow().inWholeNanoseconds
+        nanos[phase] = (nanos[phase] ?: 0L) + ns
+        calls[phase] = (calls[phase] ?: 0L) + 1L
+        return r
+    }
+
+    public fun reset() {
+        nanos.clear()
+        calls.clear()
+    }
+
+    /** Buckets sorted by accumulated time, descending, with % of the bucket total. */
+    public fun report(): String {
+        val total = nanos.values.sum()
+        fun ms(ns: Long) = ns / 1_000_000.0
+        fun pct(ns: Long) = if (total > 0) 100.0 * ns / total else 0.0
+        val width = (nanos.keys.maxOfOrNull { it.length } ?: 8).coerceAtLeast(8)
+        return buildString {
+            appendLine("[PhaseProfile] decode phase breakdown (buckets overlap matmul time; see KernelProfile):")
+            for ((phase, ns) in nanos.entries.sortedByDescending { it.value }) {
+                appendLine(
+                    "  ${phase.padEnd(width)} : ${ms(ns)} ms over ${calls[phase]} calls (${pct(ns)}%)"
+                )
+            }
+            append("  ${"phase total".padEnd(width)} : ${ms(total)} ms")
+        }
+    }
+}

transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/ResidualAdd.kt

@@ -30,6 +30,6 @@ public class ResidualAdd<T : DType, V>(
     override fun onForward(input: Tensor<T, V>, ctx: ExecutionContext): Tensor<T, V> {
         val skip = savedInput ?: return input
         savedInput = null // consume
-        return ctx.ops.add(input, skip)
+        return PhaseProfile.time("residual") { ctx.ops.add(input, skip) }
     }
 }