feat(memory): KvCacheStore declares its Format (dtype + encoding)

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

@@ -1250,6 +1250,10 @@ public final class sk/ainet/lang/memory/plan/MemoryPlan {
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class sk/ainet/lang/memory/plan/MemoryPlanKt {
+	public static final fun kvBytesFor (Lsk/ainet/lang/memory/Format;J)J
+}
+
 public final class sk/ainet/lang/memory/plan/MemoryPlans {
 	public static final field HEAP_HEADROOM_BYTES J
 	public static final field INSTANCE Lsk/ainet/lang/memory/plan/MemoryPlans;
@@ -6314,13 +6318,17 @@ public final class sk/ainet/lang/tensor/storage/DefaultKvCacheStore : sk/ainet/l
 	public fun evict (I)V
 	public fun getCurrentSeqLen ()I
 	public fun getHeadDim ()I
+	public fun getKeyBytesPerElement ()D
 	public fun getKeyEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getKeyFormat ()Lsk/ainet/lang/memory/Format;
 	public fun getMaxSeqLen ()I
 	public fun getNumHeads ()I
 	public fun getNumLayers ()I
 	public fun getPlacement ()Lsk/ainet/lang/tensor/storage/Placement;
 	public final fun getPreallocatedBytes ()J
+	public fun getValueBytesPerElement ()D
 	public fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getValueFormat ()Lsk/ainet/lang/memory/Format;
 	public fun memoryReport ()Lsk/ainet/lang/tensor/storage/KvCacheMemoryReport;
 	public fun readKeyStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage;
 	public fun readKeys (III)[F
@@ -6358,24 +6366,28 @@ public abstract interface annotation class sk/ainet/lang/tensor/storage/KvCacheB
 
 public final class sk/ainet/lang/tensor/storage/KvCacheConfig {
 	public static final field Companion Lsk/ainet/lang/tensor/storage/KvCacheConfig$Companion;
-	public fun <init> (IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;)V
-	public synthetic fun <init> (IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)V
+	public synthetic fun <init> (IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun component1 ()I
 	public final fun component2 ()I
 	public final fun component3 ()I
 	public final fun component4 ()I
 	public final fun component5 ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
 	public final fun component6 ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
 	public final fun component7 ()Lsk/ainet/lang/tensor/storage/Placement;
-	public final fun copy (IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/KvCacheConfig;
-	public static synthetic fun copy$default (Lsk/ainet/lang/tensor/storage/KvCacheConfig;IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/KvCacheConfig;
+	public final fun component8 ()Lsk/ainet/lang/types/DType;
+	public final fun component9 ()Lsk/ainet/lang/types/DType;
+	public final fun copy (IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/tensor/storage/KvCacheConfig;
+	public static synthetic fun copy$default (Lsk/ainet/lang/tensor/storage/KvCacheConfig;IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/KvCacheConfig;
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun getHeadDim ()I
+	public final fun getKeyDType ()Lsk/ainet/lang/types/DType;
 	public final fun getKeyEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
 	public final fun getMaxSeqLen ()I
 	public final fun getNumHeads ()I
 	public final fun getNumLayers ()I
 	public final fun getPlacement ()Lsk/ainet/lang/tensor/storage/Placement;
+	public final fun getValueDType ()Lsk/ainet/lang/types/DType;
 	public final fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;
@@ -6430,12 +6442,16 @@ public abstract interface class sk/ainet/lang/tensor/storage/KvCacheStore {
 	public abstract fun evict (I)V
 	public abstract fun getCurrentSeqLen ()I
 	public abstract fun getHeadDim ()I
+	public fun getKeyBytesPerElement ()D
 	public abstract fun getKeyEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getKeyFormat ()Lsk/ainet/lang/memory/Format;
 	public abstract fun getMaxSeqLen ()I
 	public abstract fun getNumHeads ()I
 	public abstract fun getNumLayers ()I
 	public abstract fun getPlacement ()Lsk/ainet/lang/tensor/storage/Placement;
+	public fun getValueBytesPerElement ()D
 	public abstract fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getValueFormat ()Lsk/ainet/lang/memory/Format;
 	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;
@@ -6456,6 +6472,10 @@ public final class sk/ainet/lang/tensor/storage/KvCacheStore$Companion {
 }
 
 public final class sk/ainet/lang/tensor/storage/KvCacheStore$DefaultImpls {
+	public static fun getKeyBytesPerElement (Lsk/ainet/lang/tensor/storage/KvCacheStore;)D
+	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 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;
@@ -6929,12 +6949,16 @@ public final class sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore : sk/aine
 	public fun evict (I)V
 	public fun getCurrentSeqLen ()I
 	public fun getHeadDim ()I
+	public fun getKeyBytesPerElement ()D
 	public fun getKeyEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getKeyFormat ()Lsk/ainet/lang/memory/Format;
 	public fun getMaxSeqLen ()I
 	public fun getNumHeads ()I
 	public fun getNumLayers ()I
 	public fun getPlacement ()Lsk/ainet/lang/tensor/storage/Placement;
+	public fun getValueBytesPerElement ()D
 	public fun getValueEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getValueFormat ()Lsk/ainet/lang/memory/Format;
 	public fun memoryReport ()Lsk/ainet/lang/tensor/storage/KvCacheMemoryReport;
 	public fun readKeyStorage (III)Lsk/ainet/lang/tensor/storage/TensorStorage;
 	public fun readKeys (III)[F

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

@@ -68,6 +68,15 @@ public enum class KvCacheMode(public val label: String) {
 }
 
 /** What to plan for: the model (header only), the context length and the prefill chunk. */
+/**
+ * The KV byte width taken from a store's declared `Format` — what the planner should use instead of
+ * guessing a [KvCacheMode] (#1077). A dense FP32 ring reports 4 bytes per element, a bf16 ring 2, a
+ * TurboQuant ring its packed width.
+ */
+@ExperimentalMemoryApi
+public fun kvBytesFor(format: Format, elements: Long): Long =
+    format.physicalBytes(elements) ?: (format.dtype.sizeInBytes.toLong() * elements)
+
 @ExperimentalMemoryApi
 public data class PlanInput(
     val modelName: String,

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt

@@ -36,6 +36,13 @@ public class DefaultKvCacheStore @kotlin.jvm.JvmOverloads constructor(
     override val maxSeqLen: Int get() = config.maxSeqLen
     override val keyEncoding: TensorEncoding get() = config.keyEncoding
     override val valueEncoding: TensorEncoding get() = config.valueEncoding
+
+    /** The dense ring holds [KvCacheConfig.keyDType] elements (FP32 unless configured otherwise) — #1077. */
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    override val keyFormat: sk.ainet.lang.memory.Format get() = sk.ainet.lang.memory.Format(config.keyDType, config.keyEncoding)
+
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    override val valueFormat: sk.ainet.lang.memory.Format get() = sk.ainet.lang.memory.Format(config.valueDType, config.valueEncoding)
     override val placement: Placement get() = config.placement
 
     private var _currentSeqLen: Int = 0

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt

@@ -20,6 +20,21 @@ import sk.ainet.lang.tensor.ops.turboquant.TurboQuantPresets
  * - Asymmetric K/V policies (e.g., Q8_0 for keys, 4-bit for values)
  * - Backend-specific fused dequant+attention paths
  */
+@OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class)
+private fun bytesPerElement(format: sk.ainet.lang.memory.Format): Double {
+    val probe = 1024L
+    val bytes = format.physicalBytes(probe) ?: return format.dtype.sizeInBytes.toDouble()
+    return bytes.toDouble() / probe
+}
+
+/** Average bytes per element of [format]; fractional for packed encodings. */
+@OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class)
+internal fun kvBytesPerElement(format: sk.ainet.lang.memory.Format): Double {
+    val probe = 1024L
+    val bytes = format.physicalBytes(probe) ?: return format.dtype.sizeInBytes.toDouble()
+    return bytes.toDouble() / probe
+}
+
 public interface KvCacheStore {
 
     /** Number of transformer layers in this cache. */
@@ -43,6 +58,33 @@ public interface KvCacheStore {
     /** Encoding used for value storage. */
     public val valueEncoding: TensorEncoding
 
+    /**
+     * What the keys *are* and how they are stored: `Format(FP32, Dense(4))` for a dense ring,
+     * `Format(BF16, Dense(2))` for a narrow-float one, `Format(FP32, TurboQuantPolar(4, 128))` for a
+     * compressed one (#1077, SKEEP-003 §0 *Format*).
+     *
+     * The default derives from [keyEncoding] and assumes FP32 — what every store held before this
+     * existed — so no implementation breaks; a store that keeps something else overrides it. The
+     * memory planner reads this instead of guessing a byte width, and guessing is what made a dense
+     * FP32 ring be planned as bf16 and understated by 2× (#1074 caught it; this makes it impossible).
+     */
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    public val keyFormat: sk.ainet.lang.memory.Format
+        get() = sk.ainet.lang.memory.Format(sk.ainet.lang.types.FP32, keyEncoding)
+
+    /** What the values are and how they are stored; see [keyFormat]. */
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    public val valueFormat: sk.ainet.lang.memory.Format
+        get() = sk.ainet.lang.memory.Format(sk.ainet.lang.types.FP32, valueEncoding)
+
+    /** Bytes one key element occupies under [keyFormat] (fractional for sub-byte encodings). */
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    public val keyBytesPerElement: Double get() = kvBytesPerElement(keyFormat)
+
+    /** Bytes one value element occupies under [valueFormat]. */
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    public val valueBytesPerElement: Double get() = kvBytesPerElement(valueFormat)
+
     /** Placement intent for the cache buffers. */
     public val placement: Placement
 
@@ -227,7 +269,14 @@ public data class KvCacheConfig(
     val maxSeqLen: Int,
     val keyEncoding: TensorEncoding = TensorEncoding.Dense(4),
     val valueEncoding: TensorEncoding = TensorEncoding.Dense(4),
-    val placement: Placement = Placement.CPU_HEAP.copy(residency = Residency.PERSISTENT)
+    val placement: Placement = Placement.CPU_HEAP.copy(residency = Residency.PERSISTENT),
+    /**
+     * The dtype the key ring stores; with [keyEncoding] it forms the store's `keyFormat` (#1077).
+     * `FP32` is what the dense store has always held; `BF16`/`FP16` halve the ring.
+     */
+    val keyDType: sk.ainet.lang.types.DType = sk.ainet.lang.types.FP32,
+    /** The dtype the value ring stores; see [keyDType]. */
+    val valueDType: sk.ainet.lang.types.DType = sk.ainet.lang.types.FP32,
 ) {
     init {
         require(numLayers > 0) { "numLayers must be positive: $numLayers" }

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt

@@ -31,6 +31,13 @@ public class TurboQuantKvCacheStore(
     override val maxSeqLen: Int get() = config.maxSeqLen
     override val keyEncoding: TensorEncoding get() = config.keyEncoding
     override val valueEncoding: TensorEncoding get() = config.valueEncoding
+
+    /** Logically FP32 values held as TurboQuant codes — the encoding the config carries (#1077). */
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    override val keyFormat: sk.ainet.lang.memory.Format get() = sk.ainet.lang.memory.Format(sk.ainet.lang.types.FP32, keyEncoding)
+
+    @sk.ainet.lang.memory.ExperimentalMemoryApi
+    override val valueFormat: sk.ainet.lang.memory.Format get() = sk.ainet.lang.memory.Format(sk.ainet.lang.types.FP32, valueEncoding)
     override val placement: Placement get() = config.placement
 
     private var _currentSeqLen: Int = 0

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheFormatTest.kt

@@ -0,0 +1,58 @@
+package sk.ainet.lang.tensor.storage
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.plan.kvBytesFor
+import sk.ainet.lang.types.BF16
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * #1077: a KV store declares its `Format` (dtype **and** encoding), so the planner reads the byte
+ * width instead of guessing it — the drift that made a dense FP32 ring be planned as bf16.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+class KvCacheFormatTest {
+
+    @Test
+    fun theDenseStoreDeclaresFp32AndTheDefaultDerivesFromTheEncoding() {
+        val store = DefaultKvCacheStore(KvCacheConfig(numLayers = 2, numHeads = 2, headDim = 8, maxSeqLen = 8))
+        assertEquals(Format(FP32, TensorEncoding.Dense(4)), store.keyFormat)
+        assertEquals(Format(FP32, TensorEncoding.Dense(4)), store.valueFormat)
+        assertEquals(4.0, store.keyBytesPerElement)
+        assertEquals(4.0, store.valueBytesPerElement)
+    }
+
+    @Test
+    fun aNarrowFloatRingIsDeclarable() {
+        val config = KvCacheConfig(
+            numLayers = 2, numHeads = 2, headDim = 8, maxSeqLen = 8,
+            keyEncoding = TensorEncoding.Dense(2), valueEncoding = TensorEncoding.Dense(2),
+            keyDType = BF16, valueDType = BF16,
+        )
+        val store = DefaultKvCacheStore(config)
+        assertEquals(Format(BF16, TensorEncoding.Dense(2)), store.keyFormat)
+        assertEquals(2.0, store.keyBytesPerElement)
+        // and the planner's width follows the declaration, not a guess
+        assertEquals(2L * 1000, kvBytesFor(store.keyFormat, 1000))
+        assertEquals(4L * 1000, kvBytesFor(Format(FP32, TensorEncoding.Dense(4)), 1000))
+    }
+
+    @Test
+    fun aCompressedStoreReportsItsPackedFormat() {
+        val store = KvCacheStore.turboQuant(numLayers = 2, numHeads = 2, headDim = 8, maxSeqLen = 8)
+        assertEquals(FP32, store.keyFormat.dtype, "TurboQuant KV is logically FP32")
+        assertTrue(store.keyFormat.encoding is TensorEncoding.TurboQuantPolar || store.keyFormat.encoding is TensorEncoding.TurboQuantPolarQjl, store.keyFormat.toString())
+        assertTrue(store.keyBytesPerElement < 4.0, "a compressed ring must be narrower than FP32, was ${store.keyBytesPerElement}")
+        assertTrue(kvBytesFor(store.keyFormat, 1024) < 4L * 1024)
+    }
+
+    @Test
+    fun theStoresFormatIsWhatThePlannerShouldUse() {
+        val dense = DefaultKvCacheStore(KvCacheConfig(numLayers = 4, numHeads = 2, headDim = 16, maxSeqLen = 64))
+        val elements = 4L * 2 * 64 * 16 * 2   // layers × heads × seq × dim × (K+V)
+        assertEquals(elements * 4, kvBytesFor(dense.keyFormat, elements), "a dense ring is FP32-wide, not bf16")
+    }
+}