From 0b049c73961d6f796123bae6181b981bd6c0965c Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 22:27:57 +0200 Subject: [PATCH] feat(memory): KvCacheStore declares its Format (dtype + encoding) instead of only an encoding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1077, follow-up to the drift the M1 harness found (#1078): the planner assumed bf16 KV while DefaultKvCacheStore holds FP32, understating a dense ring by 2x. The cache could say "4 bytes per element" but not "FP32", so nothing connected the width to a dtype. - KvCacheStore.keyFormat / valueFormat: Format(dtype, encoding) — 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. Both default to Format(FP32, ), which is what every store held before, so no implementation breaks; keyBytesPerElement / valueBytesPerElement give the (possibly fractional) width. - KvCacheConfig gains keyDType / valueDType (default FP32), so a narrow-float ring is declarable; DefaultKvCacheStore reports the configured format and TurboQuantKvCacheStore its packed one. - plan.kvBytesFor(format, elements): the width the planner should use — read from the store rather than guessed through KvCacheMode. Not generics: KvCacheStore would push the storage representation onto every caller (attention, sdpa, the transformers stack) and a compressed store could not honestly satisfy V = Float. The FloatArray boundary is a decoded-value contract (SKEEP-003 rule 4) with encode/decode inside the store; Format is how the rest of the architecture already describes "what it is and how it is stored". KvCacheFormatTest: the dense store declares FP32; a narrow-float ring is declarable and the planner's width follows the declaration; a compressed store reports its packed format and a sub-FP32 width; a dense ring is planned FP32-wide, not bf16. 181/181 storage tests. Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-lang-core.api | 32 ++++++++-- .../sk/ainet/lang/memory/plan/MemoryPlan.kt | 9 +++ .../tensor/storage/DefaultKvCacheStore.kt | 7 +++ .../ainet/lang/tensor/storage/KvCacheStore.kt | 51 +++++++++++++++- .../tensor/storage/TurboQuantKvCacheStore.kt | 7 +++ .../lang/tensor/storage/KvCacheFormatTest.kt | 58 +++++++++++++++++++ 6 files changed, 159 insertions(+), 5 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheFormatTest.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 ea2d62195..2d30faf8c 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 @@ -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,8 +6366,8 @@ 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 (IIIILsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/Placement;)V - public synthetic fun (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 (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 (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 @@ -6367,15 +6375,19 @@ public final class sk/ainet/lang/tensor/storage/KvCacheConfig { 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 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt index 1c3a6acb6..f5511e06e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt +++ b/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, 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 7f71b2203..4f12a6ced 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 @@ -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 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt index a0cebebcc..6a75d8e4c 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/KvCacheStore.kt +++ b/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" } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt index 972d49901..ba3ed136b 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt +++ b/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 diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheFormatTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/KvCacheFormatTest.kt new file mode 100644 index 000000000..5839ac5a3 --- /dev/null +++ b/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") + } +}