diff --git a/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt b/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt index b1d425caf..bb422504f 100644 --- a/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt +++ b/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt @@ -72,6 +72,34 @@ abstract class GenerateKernelMatrixTask : DefaultTask() { } appendLine("|===") appendLine("") + if (module.mapped.isNotEmpty()) { + appendLine("== Mapped serving (row-major, off-heap)") + appendLine("") + appendLine( + "Kernels that read the weight in canonical row-major GGUF file order straight " + + "from off-heap bytes (mmap'd pages or a direct buffer) — no heap staging, no " + + "prepack copy (#1189). This is what lets a model larger than the managed-heap " + + "cap decode on Android. Dense `Float32` tensors are also served from the " + + "mapping (as element views, not a matmul kernel) — see " + + "`StorageCapabilities.mappedServableEncodings` for the authoritative set the " + + "memory plan budgets against the page cache. An empty cell means the format " + + "heap-stages under `WeightResidency.MAPPED` on that platform.", + ) + appendLine("") + appendLine("[cols=\"$colSpec\", options=\"header\"]") + appendLine("|===") + append("| Weight format ") + module.platforms.forEach { append("| $it ") } + appendLine("") + appendLine("") + module.mapped.forEach { fmt -> + append("| `${fmt.name}` ") + module.platforms.forEach { p -> append("| ${fmt.byPlatform[p] ?: "—"} ") } + appendLine("") + } + appendLine("|===") + appendLine("") + } appendLine( "See also the eager backends & kernels mindmap " + "(xref:explanation/eager-execution.adoc[]) for the narrative overview and gaps.", diff --git a/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt b/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt index ad61aeb87..1084af35e 100644 --- a/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt +++ b/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt @@ -16,6 +16,12 @@ data class KernelSupportModule( val inputDtype: String = "Float32", val platforms: List = emptyList(), val formats: List = emptyList(), + /** + * Mapped serving (#1189): formats whose weight a kernel reads in canonical row-major GGUF + * file order straight from off-heap (mmap'd/direct-buffer) bytes — no heap staging, no + * prepack. Absent/empty in pre-#1189 JSON, so old files stay decodable. + */ + val mapped: List = emptyList(), ) @Serializable diff --git a/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc b/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc index 70400c938..b6d4f1f1f 100644 --- a/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc +++ b/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc @@ -20,4 +20,16 @@ Each cell is the best (highest-priority) provider that serves `Float32 × format | `Q5_0` | native-ffm | native-jni | native-cinterop | native-cinterop | scalar |=== +== Mapped serving (row-major, off-heap) + +Kernels that read the weight in canonical row-major GGUF file order straight from off-heap bytes (mmap'd pages or a direct buffer) — no heap staging, no prepack copy (#1189). This is what lets a model larger than the managed-heap cap decode on Android. Dense `Float32` tensors are also served from the mapping (as element views, not a matmul kernel) — see `StorageCapabilities.mappedServableEncodings` for the authoritative set the memory plan budgets against the page cache. An empty cell means the format heap-stages under `WeightResidency.MAPPED` on that platform. + +[cols="1,1,1,1,1,1", options="header"] +|=== +| Weight format | JVM | Android | Native·linux | Native·apple | JS/WASM + +| `Q4_K` | — | native-jni-direct | — | — | — +| `Q6_K` | — | native-jni-direct | — | — | — +|=== + See also the eager backends & kernels mindmap (xref:explanation/eager-execution.adoc[]) for the narrative overview and gaps. diff --git a/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt index bcb2be72f..b79dcf28d 100644 --- a/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt @@ -14,6 +14,7 @@ set(SKAINET_KERNEL_SOURCES ${SKAINET_KERNELS_ROOT}/src/q5_0_matmul.c ${SKAINET_KERNELS_ROOT}/src/q5_1_matmul.c ${SKAINET_KERNELS_ROOT}/src/q8_0_matmul.c + ${SKAINET_KERNELS_ROOT}/src/skainet_row_threads.c ${SKAINET_KERNELS_ROOT}/src/q4k_matmul.c ${SKAINET_KERNELS_ROOT}/src/q5k_matmul.c ${SKAINET_KERNELS_ROOT}/src/q6k_matmul.c diff --git a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c index ba9902776..2b9c58e44 100644 --- a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c +++ b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c @@ -143,6 +143,57 @@ Java_sk_ainet_exec_kernel_jni_JniKernels_q6kMatmul( inputDim, outputDim, out, outputOffset)) } +/* + * Direct-buffer row-major matmuls (#1189): the weight arrives as a direct + * ByteBuffer over mmap'd (or direct-allocated) bytes instead of a heap + * ByteArray, and stays in canonical GGUF row-major block order — which is + * what lets a model's quantized payloads never touch the managed heap. + * + * GetDirectBufferAddress is a JNI call, so it MUST run before the critical + * pins (no JNI calls are allowed between Get..Critical and Release..Critical). + * A NULL address (non-direct buffer) leaves the output untouched; the Kotlin + * caller guarantees directness by construction (DirectBufferStorage / + * MappedBufferStorage hand out direct buffers only). + */ +#define SKAINET_JNI_MATMUL_RM_DIRECT_BODY(CALL) \ + const uint8_t* w = \ + (const uint8_t*) (*env)->GetDirectBufferAddress(env, weight); \ + jfloat* in = w ? (*env)->GetPrimitiveArrayCritical(env, input, NULL) : NULL; \ + jfloat* out = in ? (*env)->GetPrimitiveArrayCritical(env, output, NULL) : NULL; \ + if (out) { \ + CALL; \ + } \ + if (out) (*env)->ReleasePrimitiveArrayCritical(env, output, out, 0); \ + if (in) (*env)->ReleasePrimitiveArrayCritical(env, input, in, JNI_ABORT); + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q4kMatmulRmDirect( + JNIEnv* env, jobject thiz, + jfloatArray input, jint inputOffset, + jobject weight, jint weightByteOffset, + jint inputDim, jint outputDim, + jfloatArray output, jint outputOffset +) { + (void) thiz; + SKAINET_JNI_MATMUL_RM_DIRECT_BODY( + skainet_q4k_matmul_rm(in, inputOffset, w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q6kMatmulRmDirect( + JNIEnv* env, jobject thiz, + jfloatArray input, jint inputOffset, + jobject weight, jint weightByteOffset, + jint inputDim, jint outputDim, + jfloatArray output, jint outputOffset +) { + (void) thiz; + SKAINET_JNI_MATMUL_RM_DIRECT_BODY( + skainet_q6k_matmul_rm(in, inputOffset, w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + /* * bitnet_gemv (SKEEP-003 §5.3, #1041): int8 activations against ternary TQ2_0 * weights. Its activation is a *byte* array, not floats, so it does not fit diff --git a/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/harness/android/M2A5DeviceMeasurement.kt b/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/harness/android/M2A5DeviceMeasurement.kt index 69f00ce08..44145bf76 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/harness/android/M2A5DeviceMeasurement.kt +++ b/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/harness/android/M2A5DeviceMeasurement.kt @@ -13,6 +13,7 @@ import sk.ainet.backend.api.kernel.KernelDispatch import sk.ainet.backend.api.kernel.KernelPacks import sk.ainet.context.DirectCpuExecutionContext import sk.ainet.exec.kernel.jni.JniKernelProvider +import sk.ainet.exec.kernel.jni.JniMappedKernelPack import sk.ainet.io.MappedRandomAccessSource import sk.ainet.io.gguf.StreamingGGUFReader import sk.ainet.io.gguf.StreamingGgufParametersLoader @@ -104,6 +105,10 @@ class M2A5DeviceMeasurement { val prepack = args.getString("prepack")?.toBoolean() ?: false val steps = arg("steps", 16) val warmup = arg("warmup", 4) + // `residency=heap` (default mapped) restores heap staging — the #1193 A/B lever for + // models that fit the cap. The plan below prices the same form the load uses (#1190). + val residency = if (args.getString("residency") == "heap") WeightResidency.HEAP else WeightResidency.MAPPED + val loadForm = WeightForm(shape = WeightShapeOrientation.OUT_IN, residency = residency) val report = StringBuilder() fun line(s: String = "") { report.append(s).append('\n') } @@ -115,13 +120,13 @@ class M2A5DeviceMeasurement { line("- device: ${Build.MANUFACTURER} ${Build.MODEL}, Android ${Build.VERSION.RELEASE} (SDK ${Build.VERSION.SDK_INT}), ABI ${Build.SUPPORTED_ABIS.firstOrNull()}") line("- ART heap cap (Runtime.maxMemory): ${mb(heapCap)}") line("- model: ${modelFile.name}, ${mb(modelFile.length())} on disk") - line("- ctx=$ctxLen, decode steps=$steps (warm-up $warmup), prepack=$prepack") + line("- ctx=$ctxLen, decode steps=$steps (warm-up $warmup), prepack=$prepack, residency=${residency.name.lowercase()}") line() // ---- plan, from the header only -------------------------------------------------- val (plan, geometry) = MappedRandomAccessSource.open(modelPath).let { src -> StreamingGGUFReader.open(src).use { reader -> - val input = reader.planInput(ctx = ctxLen) + val input = reader.planInput(ctx = ctxLen, formFor = { loadForm }) MemoryPlans.plan(input, Budget.of(heapCap)) to input.geometry } } @@ -131,6 +136,10 @@ class M2A5DeviceMeasurement { KernelDispatch.clearForTesting() runCatching { KernelPacks.install(JniKernelProvider); KernelPacks.installPacked(JniKernelProvider) } .onFailure { line("_note: JNI kernel pack unavailable (${it.message}); reference kernels serve — memory numbers stay valid, timings do not._") } + // #1189: row-major direct-buffer kernels — these serve the mapped packed weights below + // with zero copies and no prepack. + runCatching { JniMappedKernelPack.install() } + .onFailure { line("_note: mapped kernel pack unavailable (${it.message}); mapped packed weights fall back to the decoding reference._") } val sink = RecordingTraceSink() val ctx = DirectCpuExecutionContext() @@ -140,7 +149,7 @@ class M2A5DeviceMeasurement { runBlocking { StreamingGgufParametersLoader( sourceProvider = { MappedRandomAccessSource.open(modelPath) }, - weightForm = WeightForm(shape = WeightShapeOrientation.OUT_IN, residency = WeightResidency.MAPPED), + weightForm = loadForm, traceSink = sink, ).load(ctx, FP32::class) { name, t -> tensors[name] = t } } @@ -148,20 +157,29 @@ class M2A5DeviceMeasurement { val sLoaded = MemoryProbe.sample() // Account the materialized weights as model-scope allocations so plan-vs-actual sees - // them; the loader's own storages are not all sink-wired. + // them; the loader's own storages are not all sink-wired. #1189 note: physicalBytes + // (not packedData.size) — a mapped packed tensor keeps its bytes off-heap and has no + // heap ByteArray at all; those bytes are counted separately as mapped. var weightHeapBytes = 0L + var weightMappedBytes = 0L + var mappedPackedCount = 0 var syntheticStorageId = -1_000_000L for ((name, t) in tensors) { - val bytes = (t.data as? PackedBlockStorage)?.packedData?.size?.toLong() - ?: (t.shape.volume.toLong() * 4L) - weightHeapBytes += bytes + val d = t.data + val bytes = (d as? PackedBlockStorage)?.physicalBytes ?: (t.shape.volume.toLong() * 4L) + if (d is sk.ainet.lang.tensor.data.BufferPackedTensorData) { + weightMappedBytes += bytes; mappedPackedCount += 1 + } else { + weightHeapBytes += bytes + } syntheticStorageId -= 1 sink.emit(TraceEvent.Allocation(syntheticStorageId, ScopeKind.MODEL, bytes, null, "m2a5:$name")) } line("## Load") line() line("- wall time: ${loadMs} ms for ${tensors.size} tensors") - line("- materialized bytes (heap tensors + packed payloads): ${mb(weightHeapBytes)}") + line("- heap bytes (dense tensors + heap packed payloads): ${mb(weightHeapBytes)}") + line("- mapped packed payloads (#1189, off-heap): ${mb(weightMappedBytes)} in $mappedPackedCount tensors") line("- RSS before → after load: ${mb(sBefore.rssBytes)} → ${mb(sLoaded.rssBytes)} (Δ ${mb((sLoaded.rssBytes ?: 0L) - (sBefore.rssBytes ?: 0L))})") line("- major faults during load: ${sLoaded.majorFaultsSince(sBefore) ?: "—"}") line() diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBufferPackedKernels.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBufferPackedKernels.kt new file mode 100644 index 000000000..ecefa3e6d --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBufferPackedKernels.kt @@ -0,0 +1,121 @@ +package sk.ainet.exec.kernel.jni + +import java.nio.ByteBuffer +import sk.ainet.backend.api.kernel.KernelDispatch +import sk.ainet.backend.api.kernel.KernelKey +import sk.ainet.backend.api.kernel.LayoutClass +import sk.ainet.backend.api.kernel.OperandKey +import sk.ainet.backend.api.kernel.ReferenceMatmulKernel +import sk.ainet.backend.api.kernel.ViewKernel +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.DirectBufferStorage +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.MappedBufferStorage +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 + +/** + * A [ViewKernel] over the JNI **direct-buffer row-major** matmuls (#1189): the weight is a + * `BLOCKED_ROW_MAJOR` packed view whose storage is off-heap ([MappedBufferStorage] — mmap'd GGUF + * pages — or [DirectBufferStorage]), read by the native kernel straight at its address. No heap + * `ByteArray`, no relayout: this is what lets a model bigger than the ART cap decode on Android. + * + * Contrast with [sk.ainet.backend.api.kernel.PackedViewMatmulKernel], which serves the same + * encodings from heap bytes in `BLOCKED_INPUT_MAJOR` (prepacked feed) order. Registering both + * lets the dispatcher pick by what the weight actually is — canonical mapped weights match this + * key exactly and are served with zero copies. + */ +@ExperimentalMemoryApi +public class JniBufferPackedMatmulKernel( + encodingName: String, + override val key: KernelKey, + private val matmul: ( + input: FloatArray, inputOffset: Int, + weight: ByteBuffer, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) -> Unit, +) : ViewKernel { + + override val name: String = "jni-buffer-$encodingName" + + override fun run(inputs: List, out: TensorView) { + require(inputs.size == 2) { "matmul takes two operands" } + val a = inputs[0] + val w = inputs[1] + val rows = a.shape[0] + val k = a.shape[1] + val n = w.shape[0] + require(w.shape[1] == k) { "inner dimensions disagree: [$rows, $k] × [$n, ${w.shape[1]}]" } + check(w.layout.blockOrder == BlockOrder.ROW_MAJOR) { + "$name reads canonical row-major weights; this one is ${w.layout.blockOrder}" + } + + val buffer = when (val s = w.storage) { + is MappedBufferStorage -> s.buffer() + is DirectBufferStorage -> s.buffer() + else -> return fallback(inputs, out) + } + val aHeap = a.storage as? Storage.Heap ?: return fallback(inputs, out) + val oHeap = out.storage as? Storage.Heap ?: return fallback(inputs, out) + val activation = aHeap.floats ?: return fallback(inputs, out) + val output = oHeap.floats ?: return fallback(inputs, out) + // The JNI kernel takes a contiguous activation row; a strided one would be mis-indexed. + if (!a.isContiguous) return fallback(inputs, out) + + val weightOffset = (w.layout.offsetElements * w.layout.elementBytes).toInt() + for (r in 0 until rows) { + matmul( + activation, aHeap.arrayOffset + (a.layout.offsetElements + r.toLong() * k).toInt(), + buffer, weightOffset, + k, n, + output, oHeap.arrayOffset + (out.layout.offsetElements + r.toLong() * n).toInt(), + ) + } + } + + private fun fallback(inputs: List, out: TensorView) { + ReferenceMatmulKernel(key).run(inputs, out) + } + + public companion object { + /** The key this kernel serves: dense contiguous FP32 activation × row-major packed weight. */ + public fun keyFor(encoding: TensorEncoding): KernelKey = KernelKey( + op = "matmul", + operands = listOf( + OperandKey.contiguous(Format.dense(FP32)), + OperandKey(Format(FP32, encoding), LayoutClass.BLOCKED_ROW_MAJOR), + ), + ) + } +} + +/** + * Registers the direct-buffer row-major kernels (#1189) into [KernelDispatch]. Call next to + * `KernelPacks.install(JniKernelProvider)` on Android when weights are loaded with + * `WeightResidency.MAPPED` — without it a mapped packed weight falls back to the decoding + * reference kernel (correct, hours-scale slow). + */ +@ExperimentalMemoryApi +public object JniMappedKernelPack { + public fun install() { + if (!JniKernelProvider.isAvailable()) return + KernelDispatch.register( + JniBufferPackedMatmulKernel( + TensorEncoding.Q4_K.name, + JniBufferPackedMatmulKernel.keyFor(TensorEncoding.Q4_K), + JniKernels::q4kMatmulRmDirect, + ), + ) + KernelDispatch.register( + JniBufferPackedMatmulKernel( + TensorEncoding.Q6_K.name, + JniBufferPackedMatmulKernel.keyFor(TensorEncoding.Q6_K), + JniKernels::q6kMatmulRmDirect, + ), + ) + } +} diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt index 60aa085af..9101fc36c 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt @@ -171,4 +171,25 @@ public object JniKernels { inputDim: Int, outputDim: Int, output: FloatArray, outputOffset: Int, ) + + /** + * Row-major Q4_K matmul over a **direct** [java.nio.ByteBuffer] weight (#1189): the blocks + * stay in canonical GGUF file order — `(o * blocksPerRow + b) * 144` from [weightByteOffset] + * — and are read straight out of the buffer's off-heap (typically mmap'd) memory. The buffer + * MUST be direct; a heap buffer has no native address and the call writes nothing. + */ + public external fun q4kMatmulRmDirect( + input: FloatArray, inputOffset: Int, + weight: java.nio.ByteBuffer, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) + + /** Row-major Q6_K matmul over a direct ByteBuffer weight (#1189); see [q4kMatmulRmDirect]. */ + public external fun q6kMatmulRmDirect( + input: FloatArray, inputOffset: Int, + weight: java.nio.ByteBuffer, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) } diff --git a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt index e97b9d0b0..d6780f0c6 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt @@ -12,6 +12,7 @@ endif() set(SKAINET_KERNEL_SOURCES src/skainet_smoke.c src/skainet_cpu_features.c + src/skainet_row_threads.c src/q4k_matmul.c src/q5k_matmul.c src/q6k_matmul.c diff --git a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h index 43637691c..02ac35694 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h +++ b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h @@ -48,6 +48,11 @@ SKAINET_API void skainet_smoke_double(const float* input, float* output, int32_t * * Caller owns input/weight/output memory; the kernel does not retain * pointers past return. input_dim must be a multiple of 256. + * + * Threads over output rows (disjoint out[] slices, pthreads, up to 4) + * when output_dim >= 512 (#1195); single-threaded below that and on + * MSVC. Results are bit-identical either way — per output row the + * accumulation order over blocks never changes. */ SKAINET_API void skainet_q4k_matmul( const float* input, @@ -60,6 +65,29 @@ SKAINET_API void skainet_q4k_matmul( int32_t output_offset ); +/* + * Q4_K matrix-vector multiply over a ROW-MAJOR (canonical GGUF file order) + * weight (#1189). Same math and block format as skainet_q4k_matmul; the + * only difference is the weight addressing: + * weight + weight_byte_offset + (o * blocks_per_row + block_idx) * 144 + * i.e. the bytes exactly as they sit in a .gguf file — which is what lets + * mmap'd weights be fed to this kernel with no relayout copy (each row's + * blocks are read strictly sequentially). + * + * Threads over output rows when output_dim >= 512 (#1195) — see + * skainet_q4k_matmul; bit-identical to the single-threaded result. + */ +SKAINET_API void skainet_q4k_matmul_rm( + const float* input, + int32_t input_offset, + const uint8_t* weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* output, + int32_t output_offset +); + /* * Q5_K matrix-vector multiply. * @@ -100,6 +128,9 @@ SKAINET_API void skainet_q5k_matmul( * weight + weight_byte_offset + (block_idx * output_dim + o) * 210 * * input_dim must be a multiple of 256. + * + * Threads over output rows when output_dim >= 512 (#1195) — see + * skainet_q4k_matmul; bit-identical to the single-threaded result. */ SKAINET_API void skainet_q6k_matmul( const float* input, @@ -112,6 +143,28 @@ SKAINET_API void skainet_q6k_matmul( int32_t output_offset ); +/* + * Q6_K matrix-vector multiply over a ROW-MAJOR (canonical GGUF file order) + * weight (#1189). Same math and block format as skainet_q6k_matmul; the + * weight is addressed + * weight + weight_byte_offset + (o * blocks_per_row + block_idx) * 210 + * — the bytes exactly as they sit in a .gguf file, so an mmap'd weight + * needs no relayout copy. + * + * Threads over output rows when output_dim >= 512 (#1195) — see + * skainet_q4k_matmul; bit-identical to the single-threaded result. + */ +SKAINET_API void skainet_q6k_matmul_rm( + const float* input, + int32_t input_offset, + const uint8_t* weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* output, + int32_t output_offset +); + /* * Row-major FP32 SGEMM: C(m, n) = A(m, k) * B(k, n). * diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/q4k_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/q4k_matmul.c index 76d70a6b4..e5fc011f3 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/src/q4k_matmul.c +++ b/skainet-backends/skainet-backend-native-cpu/native/src/q4k_matmul.c @@ -1,6 +1,7 @@ #include "skainet_kernels.h" #include "skainet_simd.h" #include "skainet_cpu_features.h" +#include "skainet_row_threads.h" #include #include @@ -176,22 +177,119 @@ static void skainet_q4k_block_dot_generic( #endif /* - * Native Q4_K matrix-vector multiply matching the - * sk.ainet.backend.api.kernel.Q4KMatmulKernel SPI contract. Single input row - * times an `outputDim x inputDim` Q4_K-packed weight laid out - * (blockIdx * outputDim + o) * 144 bytes. - * - * Fused int8 dot path (ggml-style): the input row is quantized to Q8 ONCE per - * 256-block (reused across all output rows), then each weight sub-block is an - * int8 dot-product against the Q8 activation: - * acc += d_in[b] * ( d * Σ_s scaleIdx[s]*intDot[s] - dMin * Σ_s minIdx[s]*intSum[s] ) + * One block's contribution to out[o]: + * d_in[b] * ( d * Σ_s scaleIdx[s]*intDot[s] - dMin * Σ_s minIdx[s]*intSum[s] ) * where intDot[s] = Σ q8[i]*code[i] and intSum[s] = Σ q8[i] over the sub-block. * On AArch64 with dotprod (asimddp) the inner dot uses vdotq_s32 (16 int8 MACs * per instruction); otherwise a scalar integer fallback (auto-vectorized). - * The index mapping (groups, lo/hi sub-blocks, input alignment) is identical to - * the previous float kernel, which was parity-checked against Panama. + * Shared by the feed-order and row-major entries — identical math, so the two + * orders (and any row partition, #1195) stay bit-identical per output row. */ -SKAINET_API void skainet_q4k_matmul( +static inline float skainet_q4k_block_term( + const uint8_t* SKAINET_RESTRICT block, + const int8_t* SKAINET_RESTRICT q8_block, + float di, + int use_dp +) { + int scale_idx[Q4K_SUB_BLOCKS]; + int min_idx[Q4K_SUB_BLOCKS]; + + const uint16_t d_bits = (uint16_t) block[0] | ((uint16_t) block[1] << 8); + const uint16_t d_min_bits = (uint16_t) block[2] | ((uint16_t) block[3] << 8); + const float d = skainet_half_to_float(d_bits); + const float d_min = skainet_half_to_float(d_min_bits); + + skainet_q4k_decode_scales(block + 4, scale_idx, min_idx); + + const uint8_t* qs = block + 16; + + int64_t block_scale_dot = 0; + int64_t block_min_sum = 0; + +#if defined(SKAINET_HAVE_DOTPROD) + (void) use_dp; + skainet_q4k_block_dot_dp(qs, q8_block, scale_idx, min_idx, + &block_scale_dot, &block_min_sum); +#elif defined(SKAINET_DOTPROD_DISPATCH) + if (use_dp) { + skainet_q4k_block_dot_dp(qs, q8_block, scale_idx, min_idx, + &block_scale_dot, &block_min_sum); + } else { + skainet_q4k_block_dot_generic(qs, q8_block, scale_idx, min_idx, + &block_scale_dot, &block_min_sum); + } +#else + (void) use_dp; + skainet_q4k_block_dot_generic(qs, q8_block, scale_idx, min_idx, + &block_scale_dot, &block_min_sum); +#endif + + return di * (d * (float) block_scale_dot - d_min * (float) block_min_sum); +} + +/* Everything a row-range worker needs; read-only during the parallel section. */ +typedef struct { + const uint8_t* weight_base; /* weight + weight_byte_offset */ + const int8_t* q8; + const float* d_in; + float* out_base; + int32_t blocks_per_input_dim; + int32_t output_dim; + int use_dp; /* meaningful only under SKAINET_DOTPROD_DISPATCH */ +} skainet_q4k_ctx; + +/* + * Feed-order rows [o_start, o_end). Loop order: block OUTER, output row INNER. + * The weight is packed block-major — (blockIdx * output_dim + o) * 144 — so for + * a fixed block this range's rows are one contiguous 144·(o_end−o_start) byte + * run: reads stay sequential (prefetch- and cache-line-friendly; the o-outer + * order would stride output_dim*144 per step, a cold miss per read on an + * in-order A55). out[o] accumulates across blocks in unchanged order, so the + * result is numerically identical to the o-outer form and to any partition. + */ +static void skainet_q4k_rows_feed(void* vctx, int32_t o_start, int32_t o_end) { + const skainet_q4k_ctx* c = (const skainet_q4k_ctx*) vctx; + for (int32_t o = o_start; o < o_end; ++o) c->out_base[o] = 0.0f; + for (int32_t block_idx = 0; block_idx < c->blocks_per_input_dim; ++block_idx) { + const int8_t* q8_block = c->q8 + (size_t) block_idx * Q4K_BLOCK_SIZE; + const float di = c->d_in[block_idx]; + const uint8_t* block = c->weight_base + + ((size_t) block_idx * c->output_dim + o_start) * Q4K_BYTES_PER_BLOCK; + for (int32_t o = o_start; o < o_end; ++o, block += Q4K_BYTES_PER_BLOCK) { + c->out_base[o] += skainet_q4k_block_term(block, q8_block, di, c->use_dp); + } + } +} + +/* + * Row-major rows [o_start, o_end) (#1189): canonical GGUF file order — + * (o * blocks_per_row + b) * 144 — so an mmap'd tensor is fed as-is, no + * relayout copy. o OUTER: each row's blocks are contiguous on disk, reads stay + * strictly sequential; the Q8 activation (input_dim bytes) stays hot across + * rows. Per-row accumulation order matches the feed-order worker's. + */ +static void skainet_q4k_rows_rm(void* vctx, int32_t o_start, int32_t o_end) { + const skainet_q4k_ctx* c = (const skainet_q4k_ctx*) vctx; + const uint8_t* block = c->weight_base + + (size_t) o_start * c->blocks_per_input_dim * Q4K_BYTES_PER_BLOCK; + for (int32_t o = o_start; o < o_end; ++o) { + float acc = 0.0f; + for (int32_t block_idx = 0; block_idx < c->blocks_per_input_dim; + ++block_idx, block += Q4K_BYTES_PER_BLOCK) { + acc += skainet_q4k_block_term(block, c->q8 + (size_t) block_idx * Q4K_BLOCK_SIZE, + c->d_in[block_idx], c->use_dp); + } + c->out_base[o] = acc; + } +} + +/* + * Shared entry: quantize the input row to Q8 once (reused across all rows and + * all threads — read-only after this point), then run the worker over the + * output rows, threaded per skainet_row_threads.h (#1195: ≥512 rows → up to 4 + * pthreads; below that, or on MSVC, the calling thread does all rows). + */ +static void skainet_q4k_matmul_run( const float* SKAINET_RESTRICT input, int32_t input_offset, const uint8_t* SKAINET_RESTRICT weight, @@ -199,20 +297,14 @@ SKAINET_API void skainet_q4k_matmul( int32_t input_dim, int32_t output_dim, float* SKAINET_RESTRICT output, - int32_t output_offset + int32_t output_offset, + skainet_row_range_fn worker ) { if (output_dim <= 0 || input_dim <= 0) return; -#ifdef SKAINET_DOTPROD_DISPATCH - /* One probe per matmul call; cached in skainet_cpu_has_dotprod. */ - const int use_dp = skainet_cpu_has_dotprod(); -#endif - const int32_t blocks_per_input_dim = input_dim / Q4K_BLOCK_SIZE; const float* in_base = input + input_offset; - float* out_base = output + output_offset; - /* Pre-quantize the whole input row to Q8 once (reused across all o). */ int8_t* q8 = (int8_t*) malloc((size_t) input_dim * sizeof(int8_t)); float* d_in = (float*) malloc((size_t) blocks_per_input_dim * sizeof(float)); if (q8 == NULL || d_in == NULL) { free(q8); free(d_in); return; } @@ -221,62 +313,64 @@ SKAINET_API void skainet_q4k_matmul( q8 + (size_t) b * Q4K_BLOCK_SIZE); } - int scale_idx[Q4K_SUB_BLOCKS]; - int min_idx[Q4K_SUB_BLOCKS]; - - /* - * Loop order: block OUTER, output row INNER. The weight is packed - * block-major — (blockIdx * output_dim + o) * 144 — so for a fixed block, - * consecutive `o` are exactly 144 bytes apart: the weight bytes are read - * strictly sequentially (prefetch- and cache-line-friendly). The reverse - * order (o outer) strides output_dim*144 bytes per step (~295 KB on the - * down-proj), which on an in-order A55 with small caches makes every weight - * read a cold miss and dominates runtime regardless of inner-loop compute. - * out_base[o] is accumulated across blocks (output_dim*4 bytes stays hot in - * cache); the accumulation order over blocks is unchanged, so this is - * numerically identical to the o-outer form. - */ - for (int32_t o = 0; o < output_dim; ++o) out_base[o] = 0.0f; - - for (int32_t block_idx = 0; block_idx < blocks_per_input_dim; ++block_idx) { - const int8_t* q8_block = q8 + (size_t) block_idx * Q4K_BLOCK_SIZE; - const float di = d_in[block_idx]; - const uint8_t* block = weight + weight_byte_offset - + (size_t)(block_idx * output_dim) * Q4K_BYTES_PER_BLOCK; - - for (int32_t o = 0; o < output_dim; ++o, block += Q4K_BYTES_PER_BLOCK) { - const uint16_t d_bits = (uint16_t) block[0] | ((uint16_t) block[1] << 8); - const uint16_t d_min_bits = (uint16_t) block[2] | ((uint16_t) block[3] << 8); - const float d = skainet_half_to_float(d_bits); - const float d_min = skainet_half_to_float(d_min_bits); - - skainet_q4k_decode_scales(block + 4, scale_idx, min_idx); - - const uint8_t* qs = block + 16; - - int64_t block_scale_dot = 0; - int64_t block_min_sum = 0; - -#if defined(SKAINET_HAVE_DOTPROD) - skainet_q4k_block_dot_dp(qs, q8_block, scale_idx, min_idx, - &block_scale_dot, &block_min_sum); -#elif defined(SKAINET_DOTPROD_DISPATCH) - if (use_dp) { - skainet_q4k_block_dot_dp(qs, q8_block, scale_idx, min_idx, - &block_scale_dot, &block_min_sum); - } else { - skainet_q4k_block_dot_generic(qs, q8_block, scale_idx, min_idx, - &block_scale_dot, &block_min_sum); - } + skainet_q4k_ctx ctx; + ctx.weight_base = weight + weight_byte_offset; + ctx.q8 = q8; + ctx.d_in = d_in; + ctx.out_base = output + output_offset; + ctx.blocks_per_input_dim = blocks_per_input_dim; + ctx.output_dim = output_dim; +#ifdef SKAINET_DOTPROD_DISPATCH + /* One probe per matmul call; cached in skainet_cpu_has_dotprod. */ + ctx.use_dp = skainet_cpu_has_dotprod(); #else - skainet_q4k_block_dot_generic(qs, q8_block, scale_idx, min_idx, - &block_scale_dot, &block_min_sum); + ctx.use_dp = 0; #endif - out_base[o] += di * (d * (float) block_scale_dot - d_min * (float) block_min_sum); - } - } + skainet_run_rows(worker, &ctx, output_dim); free(q8); free(d_in); } + +/* + * Native Q4_K matrix-vector multiply matching the + * sk.ainet.backend.api.kernel.Q4KMatmulKernel SPI contract. Single input row + * times an `outputDim x inputDim` Q4_K-packed weight laid out + * (blockIdx * outputDim + o) * 144 bytes. Threads over output rows when + * outputDim >= 512 (#1195); parity-checked against Panama. + */ +SKAINET_API void skainet_q4k_matmul( + const float* SKAINET_RESTRICT input, + int32_t input_offset, + const uint8_t* SKAINET_RESTRICT weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* SKAINET_RESTRICT output, + int32_t output_offset +) { + skainet_q4k_matmul_run(input, input_offset, weight, weight_byte_offset, + input_dim, output_dim, output, output_offset, + skainet_q4k_rows_feed); +} + +/* + * Row-major variant (#1189): the weight stays in canonical GGUF file order — + * (o * blocks_per_row + b) * 144 — see skainet_q4k_rows_rm. Bit-identical to + * the feed-order kernel; threads over output rows when outputDim >= 512. + */ +SKAINET_API void skainet_q4k_matmul_rm( + const float* SKAINET_RESTRICT input, + int32_t input_offset, + const uint8_t* SKAINET_RESTRICT weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* SKAINET_RESTRICT output, + int32_t output_offset +) { + skainet_q4k_matmul_run(input, input_offset, weight, weight_byte_offset, + input_dim, output_dim, output, output_offset, + skainet_q4k_rows_rm); +} diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/q6k_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/q6k_matmul.c index 2a64fce12..ea7d78742 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/src/q6k_matmul.c +++ b/skainet-backends/skainet-backend-native-cpu/native/src/q6k_matmul.c @@ -1,6 +1,7 @@ #include "skainet_kernels.h" #include "skainet_simd.h" #include "skainet_cpu_features.h" +#include "skainet_row_threads.h" #include #include @@ -152,18 +153,100 @@ static int64_t skainet_q6k_weighted_dot_generic(const int8_t* SKAINET_RESTRICT q #endif /* - * Native Q6_K matrix-vector multiply matching the - * sk.ainet.backend.api.kernel.Q6KMatmulKernel SPI contract. A single - * input row times an `outputDim x inputDim` Q6_K-packed weight tensor - * laid out (blockIdx * outputDim + o) * 210 bytes. - * - * Fused int8 dot path (ggml-style, mirrors q4k_matmul.c): the input row is - * quantized to Q8 ONCE per 256-block (reused across all output rows), the 6-bit - * weight is unpacked to centered int8 codes, and each scale-group is an int8 - * dot (vdotq_s32 on dotprod targets) — no 256-float scratch, no per-element - * float multiply. acc = d · d_in · Σ_g sc[g]·Σ_{i∈g} q8[i]·codes[i]. + * One block's contribution to out[o]: the 6-bit weight is unpacked to centered + * int8 codes and each scale-group is an int8 dot (vdotq_s32 on dotprod + * targets) — acc term = d · d_in · Σ_g sc[g]·Σ_{i∈g} q8[i]·codes[i]. + * `codes` is the caller's per-thread 256-byte scratch (#1195). Shared by the + * feed-order and row-major entries, so both orders (and any row partition) + * stay bit-identical per output row. */ -SKAINET_API void skainet_q6k_matmul( +static inline float skainet_q6k_block_term( + const uint8_t* SKAINET_RESTRICT block, + const int8_t* SKAINET_RESTRICT q8_block, + float di, + int use_dp, + int8_t* SKAINET_RESTRICT codes +) { + const uint16_t d_bits = (uint16_t) block[Q6K_D_OFFSET] + | ((uint16_t) block[Q6K_D_OFFSET + 1] << 8); + const float d = skainet_q6k_half_to_float(d_bits); + const int8_t* sc = (const int8_t*)(block + Q6K_SCALES_OFFSET); + + skainet_q6k_unpack_codes(block, codes); +#if defined(SKAINET_HAVE_DOTPROD) + (void) use_dp; + const int64_t wdot = skainet_q6k_weighted_dot_dp(q8_block, codes, sc); +#elif defined(SKAINET_DOTPROD_DISPATCH) + const int64_t wdot = use_dp + ? skainet_q6k_weighted_dot_dp(q8_block, codes, sc) + : skainet_q6k_weighted_dot_generic(q8_block, codes, sc); +#else + (void) use_dp; + const int64_t wdot = skainet_q6k_weighted_dot_generic(q8_block, codes, sc); +#endif + + return d * di * (float) wdot; +} + +/* Everything a row-range worker needs; read-only during the parallel section. */ +typedef struct { + const uint8_t* weight_base; /* weight + weight_byte_offset */ + const int8_t* q8; + const float* d_in; + float* out_base; + int32_t blocks_per_input_dim; + int32_t output_dim; + int use_dp; /* meaningful only under SKAINET_DOTPROD_DISPATCH */ +} skainet_q6k_ctx; + +/* + * Feed-order rows [o_start, o_end): block OUTER, row INNER — see q4k_matmul.c + * for the cache rationale. This range's rows are one contiguous run per block; + * out[o] accumulates across blocks in unchanged order under any partition. + */ +static void skainet_q6k_rows_feed(void* vctx, int32_t o_start, int32_t o_end) { + const skainet_q6k_ctx* c = (const skainet_q6k_ctx*) vctx; + int8_t codes[Q6K_BLOCK_SIZE]; + for (int32_t o = o_start; o < o_end; ++o) c->out_base[o] = 0.0f; + for (int32_t block_idx = 0; block_idx < c->blocks_per_input_dim; ++block_idx) { + const int8_t* q8_block = c->q8 + (size_t) block_idx * Q6K_BLOCK_SIZE; + const float di = c->d_in[block_idx]; + const uint8_t* block = c->weight_base + + ((size_t) block_idx * c->output_dim + o_start) * Q6K_BYTES_PER_BLOCK; + for (int32_t o = o_start; o < o_end; ++o, block += Q6K_BYTES_PER_BLOCK) { + c->out_base[o] += skainet_q6k_block_term(block, q8_block, di, c->use_dp, codes); + } + } +} + +/* + * Row-major rows [o_start, o_end) (#1189): canonical GGUF file order — + * (o * blocks_per_row + b) * 210 — an mmap'd tensor is fed as-is, no relayout + * copy; each row's blocks are contiguous on disk. Per-row accumulation order + * matches the feed-order worker's. + */ +static void skainet_q6k_rows_rm(void* vctx, int32_t o_start, int32_t o_end) { + const skainet_q6k_ctx* c = (const skainet_q6k_ctx*) vctx; + int8_t codes[Q6K_BLOCK_SIZE]; + const uint8_t* block = c->weight_base + + (size_t) o_start * c->blocks_per_input_dim * Q6K_BYTES_PER_BLOCK; + for (int32_t o = o_start; o < o_end; ++o) { + float acc = 0.0f; + for (int32_t block_idx = 0; block_idx < c->blocks_per_input_dim; + ++block_idx, block += Q6K_BYTES_PER_BLOCK) { + acc += skainet_q6k_block_term(block, c->q8 + (size_t) block_idx * Q6K_BLOCK_SIZE, + c->d_in[block_idx], c->use_dp, codes); + } + c->out_base[o] = acc; + } +} + +/* + * Shared entry: quantize the input row to Q8 once (read-only afterwards, shared + * by all threads), then run the worker over the output rows, threaded per + * skainet_row_threads.h (#1195). + */ +static void skainet_q6k_matmul_run( const float* SKAINET_RESTRICT input, int32_t input_offset, const uint8_t* SKAINET_RESTRICT weight, @@ -171,20 +254,14 @@ SKAINET_API void skainet_q6k_matmul( int32_t input_dim, int32_t output_dim, float* SKAINET_RESTRICT output, - int32_t output_offset + int32_t output_offset, + skainet_row_range_fn worker ) { if (output_dim <= 0 || input_dim <= 0) return; -#ifdef SKAINET_DOTPROD_DISPATCH - /* One probe per matmul call; cached in skainet_cpu_has_dotprod. */ - const int use_dp = skainet_cpu_has_dotprod(); -#endif - const int32_t blocks_per_input_dim = input_dim / Q6K_BLOCK_SIZE; const float* in_base = input + input_offset; - float* out_base = output + output_offset; - /* Pre-quantize the whole input row to Q8 once (reused across all o). */ int8_t* q8 = (int8_t*) malloc((size_t) input_dim * sizeof(int8_t)); float* d_in = (float*) malloc((size_t) blocks_per_input_dim * sizeof(float)); if (q8 == NULL || d_in == NULL) { free(q8); free(d_in); return; } @@ -193,45 +270,64 @@ SKAINET_API void skainet_q6k_matmul( q8 + (size_t) b * Q6K_BLOCK_SIZE); } - int8_t codes[Q6K_BLOCK_SIZE]; - - /* - * Loop order: block OUTER, output row INNER — see q4k_matmul.c for the - * rationale. The weight is block-major (blockIdx*output_dim + o)*210, so for - * a fixed block consecutive `o` are 210 bytes apart: the weight bytes are - * read sequentially (cache/prefetch friendly) instead of striding - * output_dim*210 per step. out_base[o] accumulates across blocks; the order - * over blocks is unchanged. - */ - for (int32_t o = 0; o < output_dim; ++o) out_base[o] = 0.0f; - - for (int32_t block_idx = 0; block_idx < blocks_per_input_dim; ++block_idx) { - const int8_t* q8_block = q8 + (size_t) block_idx * Q6K_BLOCK_SIZE; - const float di = d_in[block_idx]; - const uint8_t* block = weight + weight_byte_offset - + (size_t)(block_idx * output_dim) * Q6K_BYTES_PER_BLOCK; - - for (int32_t o = 0; o < output_dim; ++o, block += Q6K_BYTES_PER_BLOCK) { - const uint16_t d_bits = (uint16_t) block[Q6K_D_OFFSET] - | ((uint16_t) block[Q6K_D_OFFSET + 1] << 8); - const float d = skainet_q6k_half_to_float(d_bits); - const int8_t* sc = (const int8_t*)(block + Q6K_SCALES_OFFSET); - - skainet_q6k_unpack_codes(block, codes); -#if defined(SKAINET_HAVE_DOTPROD) - const int64_t wdot = skainet_q6k_weighted_dot_dp(q8_block, codes, sc); -#elif defined(SKAINET_DOTPROD_DISPATCH) - const int64_t wdot = use_dp - ? skainet_q6k_weighted_dot_dp(q8_block, codes, sc) - : skainet_q6k_weighted_dot_generic(q8_block, codes, sc); + skainet_q6k_ctx ctx; + ctx.weight_base = weight + weight_byte_offset; + ctx.q8 = q8; + ctx.d_in = d_in; + ctx.out_base = output + output_offset; + ctx.blocks_per_input_dim = blocks_per_input_dim; + ctx.output_dim = output_dim; +#ifdef SKAINET_DOTPROD_DISPATCH + /* One probe per matmul call; cached in skainet_cpu_has_dotprod. */ + ctx.use_dp = skainet_cpu_has_dotprod(); #else - const int64_t wdot = skainet_q6k_weighted_dot_generic(q8_block, codes, sc); + ctx.use_dp = 0; #endif - out_base[o] += d * di * (float) wdot; - } - } + skainet_run_rows(worker, &ctx, output_dim); free(q8); free(d_in); } + +/* + * Native Q6_K matrix-vector multiply matching the + * sk.ainet.backend.api.kernel.Q6KMatmulKernel SPI contract. A single + * input row times an `outputDim x inputDim` Q6_K-packed weight tensor + * laid out (blockIdx * outputDim + o) * 210 bytes. Threads over output + * rows when outputDim >= 512 (#1195). + */ +SKAINET_API void skainet_q6k_matmul( + const float* SKAINET_RESTRICT input, + int32_t input_offset, + const uint8_t* SKAINET_RESTRICT weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* SKAINET_RESTRICT output, + int32_t output_offset +) { + skainet_q6k_matmul_run(input, input_offset, weight, weight_byte_offset, + input_dim, output_dim, output, output_offset, + skainet_q6k_rows_feed); +} + +/* + * Row-major variant (#1189): the weight stays in canonical GGUF file order — + * (o * blocks_per_row + b) * 210 — see skainet_q6k_rows_rm. Bit-identical to + * the feed-order kernel; threads over output rows when outputDim >= 512. + */ +SKAINET_API void skainet_q6k_matmul_rm( + const float* SKAINET_RESTRICT input, + int32_t input_offset, + const uint8_t* SKAINET_RESTRICT weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* SKAINET_RESTRICT output, + int32_t output_offset +) { + skainet_q6k_matmul_run(input, input_offset, weight, weight_byte_offset, + input_dim, output_dim, output, output_offset, + skainet_q6k_rows_rm); +} diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/skainet_row_threads.c b/skainet-backends/skainet-backend-native-cpu/native/src/skainet_row_threads.c new file mode 100644 index 000000000..10a319ca8 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/skainet_row_threads.c @@ -0,0 +1,231 @@ +/* + * Row-range threading shared by the packed matmul kernels (#1195). + * + * The packed kernels parallelize over OUTPUT ROWS: participants pull disjoint + * `[o_start, o_end)` grains of the output vector, so there are no + * accumulation races and no synchronization beyond the completion barrier. + * Per output row the accumulation order over input blocks is unchanged, and + * a row's result does not depend on which thread computes it — threaded + * results are bit-identical to single-threaded ones (the parity suites are + * the oracle for that claim). + * + * Execution model, shaped by a series of Pixel 8a measurements (#1195, all + * Qwen2.5-1.5B Q4_K_M decode steps; single-threaded baseline 153 ms): + * + * 1. POOL, NOT CREATE/JOIN. The obvious create/join-per-call variant ran + * 994 ms/step: ~600 pthread_create per step against cores in deep cpuidle + * pays milliseconds of wakeup latency each. Workers are created once, + * lazily (pthread_once), and live for the process. + * + * 2. SPIN BRIEFLY, THEN PARK. A pool whose workers sleep between jobs still + * measured only 115–306 ms/step across chunking variants — sub-millisecond + * parallel bursts separated by sleeps never accumulate per-thread + * utilization, so EAS/schedutil keeps the workers on little cores at low + * clocks while the single-threaded caller would have pegged one big core + * at max clock. Workers therefore spin (`yield`) on the job epoch for + * ~SKAINET_SPIN_ITERS before parking on the condvar: during a decode the + * gaps between matmuls are far shorter than the spin window, utilization + * stays pegged, and the scheduler answers with big cores and full clocks + * (the same reason llama.cpp's thread pool spins). Once work stops + * arriving, everyone parks — no battery burn at idle. + * + * 3. GUIDED grains — remaining/(2·parts), floored at SKAINET_MATMUL_GRAIN — + * off an atomic cursor: long contiguous streams first (prefetch-friendly), + * shrinking toward the tail so a straggler holds a small tail rather than + * a quarter of the matrix. On symmetric cores this degrades to an even + * split; nothing here is tuned to one SoC's topology. + * + * Threading engages only when `n` reaches the threshold — below it fixed + * costs dominate any win, and tiny projections (e.g. GQA k/v with output_dim + * 256) stay single-threaded on purpose. A concurrent second caller while the + * pool is busy simply runs its rows on its own thread (correct, unshared); + * so does everything if worker creation ever failed. + * + * MSVC has no ; there the runner degrades to a plain call on the + * caller's thread. + */ +#include "skainet_row_threads.h" + +#if !defined(_MSC_VER) + +#include +#include + +#define SKAINET_POOL_WORKERS (SKAINET_MATMUL_THREADS - 1) +#define SKAINET_MATMUL_GRAIN 64 +/* ~a millisecond of `yield`s — longer than the gaps between a decode step's + * matmul calls, far shorter than "the model stopped decoding". */ +#define SKAINET_SPIN_ITERS (1u << 20) + +#if defined(__aarch64__) || defined(__arm__) +#define SKAINET_CPU_RELAX() __asm__ __volatile__("yield" ::: "memory") +#elif defined(__x86_64__) || defined(__i386__) +#define SKAINET_CPU_RELAX() __asm__ __volatile__("pause" ::: "memory") +#else +#define SKAINET_CPU_RELAX() ((void) 0) +#endif + +typedef struct { + pthread_mutex_t m; /* serializes callers; guards the park/wake handoffs */ + pthread_cond_t cv_work; + pthread_cond_t cv_done; + skainet_row_range_fn fn; /* job fields: written before the epoch release-store */ + void* ctx; + int32_t n; + int parts; + atomic_int_fast32_t cursor; /* next unclaimed row of the current job */ + atomic_ulong epoch; /* release-published per job; the workers' work signal */ + atomic_int remaining; /* workers not yet finished with the current job */ + atomic_int work_waiters; /* workers parked on cv_work (broadcast only then) */ + atomic_int done_waiter; /* caller parked on cv_done (signal only then) */ + int busy; /* a job is in flight (second callers go solo) */ + int workers_alive; +} skainet_row_pool; + +static skainet_row_pool skainet_g_row_pool = { + PTHREAD_MUTEX_INITIALIZER, PTHREAD_COND_INITIALIZER, PTHREAD_COND_INITIALIZER, + NULL, NULL, 0, 0, 0, 0UL, 0, 0, 0, 0, 0, +}; +static pthread_once_t skainet_g_row_pool_once = PTHREAD_ONCE_INIT; + +/* Guided sizing (OpenMP `schedule(guided)` shape): see file header, point 3. */ +static void skainet_row_pool_drain(skainet_row_range_fn fn, void* ctx, int32_t n, int parts) { + int_fast32_t cur = atomic_load_explicit(&skainet_g_row_pool.cursor, memory_order_relaxed); + for (;;) { + if ((int32_t) cur >= n) return; + int32_t want = (n - (int32_t) cur) / (2 * parts); + if (want < SKAINET_MATMUL_GRAIN) want = SKAINET_MATMUL_GRAIN; + if (atomic_compare_exchange_weak_explicit( + &skainet_g_row_pool.cursor, &cur, cur + want, + memory_order_relaxed, memory_order_relaxed)) { + const int32_t s = (int32_t) cur; + int32_t e = s + want; + if (e > n) e = n; + fn(ctx, s, e); + cur = atomic_load_explicit(&skainet_g_row_pool.cursor, memory_order_relaxed); + } + /* CAS failure reloaded `cur`; loop retries with the fresh value. */ + } +} + +static void* skainet_row_pool_worker(void* arg) { + (void) arg; + unsigned long seen = 0UL; + for (;;) { + /* Spin for the next job; park only when none arrives in the window. */ + unsigned spins = 0; + while (atomic_load_explicit(&skainet_g_row_pool.epoch, memory_order_acquire) == seen) { + if (++spins >= SKAINET_SPIN_ITERS) { + pthread_mutex_lock(&skainet_g_row_pool.m); + atomic_fetch_add_explicit(&skainet_g_row_pool.work_waiters, 1, memory_order_relaxed); + while (atomic_load_explicit(&skainet_g_row_pool.epoch, memory_order_acquire) == seen) { + pthread_cond_wait(&skainet_g_row_pool.cv_work, &skainet_g_row_pool.m); + } + atomic_fetch_sub_explicit(&skainet_g_row_pool.work_waiters, 1, memory_order_relaxed); + pthread_mutex_unlock(&skainet_g_row_pool.m); + break; + } + SKAINET_CPU_RELAX(); + } + seen = atomic_load_explicit(&skainet_g_row_pool.epoch, memory_order_acquire); + + /* Job fields were written before the epoch release-store — the acquire + * above orders these plain reads, and they are stable while busy. */ + skainet_row_pool_drain(skainet_g_row_pool.fn, skainet_g_row_pool.ctx, + skainet_g_row_pool.n, skainet_g_row_pool.parts); + + if (atomic_fetch_sub_explicit(&skainet_g_row_pool.remaining, 1, memory_order_acq_rel) == 1) { + /* Last one out: the park/wake handoff must be decided under the + * mutex — a bare done_waiter load could miss a caller that is + * between setting the flag and blocking, and sleep it forever. + * Cost: one lock/unlock per job, by one thread. */ + pthread_mutex_lock(&skainet_g_row_pool.m); + if (atomic_load_explicit(&skainet_g_row_pool.done_waiter, memory_order_relaxed) != 0) { + pthread_cond_signal(&skainet_g_row_pool.cv_done); + } + pthread_mutex_unlock(&skainet_g_row_pool.m); + } + } + /* unreachable */ +} + +static void skainet_row_pool_init(void) { + for (int i = 0; i < SKAINET_POOL_WORKERS; ++i) { + pthread_t t; + if (pthread_create(&t, NULL, skainet_row_pool_worker, NULL) != 0) { + break; /* fewer participants; the cursor still covers every row */ + } + pthread_detach(t); + ++skainet_g_row_pool.workers_alive; + } +} + +/* + * Run `fn` over rows [0, n): every participant (workers + the calling + * thread) pulls guided grains from the cursor; the caller then spins briefly + * on the completion count before parking. Single-threaded when n is under + * the threshold, when the pool is busy with another caller's job, or when no + * workers exist. + */ +void skainet_run_rows(skainet_row_range_fn fn, void* ctx, int32_t n) { + if (n <= 0) return; + if (n < SKAINET_MATMUL_THREAD_THRESHOLD) { + fn(ctx, 0, n); + return; + } + pthread_once(&skainet_g_row_pool_once, skainet_row_pool_init); + + pthread_mutex_lock(&skainet_g_row_pool.m); + if (skainet_g_row_pool.workers_alive == 0 || skainet_g_row_pool.busy) { + pthread_mutex_unlock(&skainet_g_row_pool.m); + fn(ctx, 0, n); + return; + } + skainet_g_row_pool.busy = 1; + skainet_g_row_pool.fn = fn; + skainet_g_row_pool.ctx = ctx; + skainet_g_row_pool.n = n; + skainet_g_row_pool.parts = skainet_g_row_pool.workers_alive + 1; + atomic_store_explicit(&skainet_g_row_pool.cursor, 0, memory_order_relaxed); + atomic_store_explicit(&skainet_g_row_pool.remaining, skainet_g_row_pool.workers_alive, + memory_order_relaxed); + /* Publish: job fields above happen-before this release-store. */ + atomic_store_explicit(&skainet_g_row_pool.epoch, + atomic_load_explicit(&skainet_g_row_pool.epoch, memory_order_relaxed) + 1UL, + memory_order_release); + if (atomic_load_explicit(&skainet_g_row_pool.work_waiters, memory_order_relaxed) != 0) { + pthread_cond_broadcast(&skainet_g_row_pool.cv_work); + } + const int parts = skainet_g_row_pool.parts; + pthread_mutex_unlock(&skainet_g_row_pool.m); + + skainet_row_pool_drain(fn, ctx, n, parts); + + /* Spin briefly for the stragglers' tail, then park. */ + unsigned spins = 0; + while (atomic_load_explicit(&skainet_g_row_pool.remaining, memory_order_acquire) > 0) { + if (++spins >= SKAINET_SPIN_ITERS) { + pthread_mutex_lock(&skainet_g_row_pool.m); + atomic_store_explicit(&skainet_g_row_pool.done_waiter, 1, memory_order_release); + while (atomic_load_explicit(&skainet_g_row_pool.remaining, memory_order_acquire) > 0) { + pthread_cond_wait(&skainet_g_row_pool.cv_done, &skainet_g_row_pool.m); + } + atomic_store_explicit(&skainet_g_row_pool.done_waiter, 0, memory_order_relaxed); + pthread_mutex_unlock(&skainet_g_row_pool.m); + break; + } + SKAINET_CPU_RELAX(); + } + + pthread_mutex_lock(&skainet_g_row_pool.m); + skainet_g_row_pool.busy = 0; + pthread_mutex_unlock(&skainet_g_row_pool.m); +} + +#else /* MSVC: no pthreads — single-threaded, same numerics */ + +void skainet_run_rows(skainet_row_range_fn fn, void* ctx, int32_t n) { + if (n > 0) fn(ctx, 0, n); +} + +#endif diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/skainet_row_threads.h b/skainet-backends/skainet-backend-native-cpu/native/src/skainet_row_threads.h new file mode 100644 index 000000000..6f78f57d6 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/skainet_row_threads.h @@ -0,0 +1,25 @@ +/* + * Row-range threading shared by the packed matmul kernels (#1195) — see + * skainet_row_threads.c for the pool and the measured rationale. + */ +#ifndef SKAINET_ROW_THREADS_H +#define SKAINET_ROW_THREADS_H + +#include + +#define SKAINET_MATMUL_THREADS 4 +#define SKAINET_MATMUL_THREAD_THRESHOLD 512 + +typedef void (*skainet_row_range_fn)(void* ctx, int32_t o_start, int32_t o_end); + +/* + * Run `fn` over rows [0, n): threaded over the shared worker pool when + * n >= SKAINET_MATMUL_THREAD_THRESHOLD, on the calling thread otherwise + * (also on MSVC, when the pool is busy with another caller, or when worker + * creation failed). Bit-identical to the single-threaded result: workers own + * disjoint [o_start, o_end) row ranges and per-row accumulation order never + * changes. + */ +void skainet_run_rows(skainet_row_range_fn fn, void* ctx, int32_t n); + +#endif /* SKAINET_ROW_THREADS_H */ diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt index c6ad891a0..d4cbc7cd6 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt @@ -56,6 +56,18 @@ class KernelSupportMatrixTest { setOf("Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K", "Q5_1", "Q5_0")), ) + /** + * Mapped serving (#1189): kernels that read the weight in canonical row-major GGUF file + * order straight from off-heap bytes (mmap/direct buffer) — the `_rm` symbols behind + * `JniBufferPackedMatmulKernel` on Android. The JVM/FFM tier is #1191; other platforms + * and formats are #1192. Must stay in lockstep with + * `StorageCapabilities.MAPPED_SERVABLE_ENCODINGS` (dense F32 is mapped there too, but as + * element-view serving, not a matmul kernel — it has no row here on purpose). + */ + private fun mappedTiers(): List = listOf( + Tier("native-jni-direct", 100, setOf("Android"), setOf("Q4_K", "Q6_K")), + ) + private fun best(fmt: String, platform: String, tiers: List): String? = tiers.filter { platform in it.platforms && fmt in it.formats }.maxByOrNull { it.priority }?.name @@ -75,6 +87,16 @@ class KernelSupportMatrixTest { .append(cells.joinToString(", ")).append("}}") sb.append(if (i == formats.lastIndex) "\n" else ",\n") } + sb.append(" ],\n") + val mapped = mappedTiers() + val mappedFormats = formats.filter { fmt -> mapped.any { fmt in it.formats } } + sb.append(" \"mapped\": [\n") + mappedFormats.forEachIndexed { i, fmt -> + val cells = platforms.mapNotNull { p -> best(fmt, p, mapped)?.let { "\"$p\": \"$it\"" } } + sb.append(" {\"name\": \"").append(fmt).append("\", \"byPlatform\": {") + .append(cells.joinToString(", ")).append("}}") + sb.append(if (i == mappedFormats.lastIndex) "\n" else ",\n") + } sb.append(" ]\n}\n") return sb.toString() } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/RowMajorMatmulParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/RowMajorMatmulParityTest.kt new file mode 100644 index 000000000..5ec8da6d6 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/RowMajorMatmulParityTest.kt @@ -0,0 +1,192 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.FunctionDescriptor +import java.lang.foreign.Linker +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.lang.invoke.MethodHandle +import kotlin.random.Random +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +/** + * Parity tests for the row-major kernel variants `skainet_q4k_matmul_rm` / + * `skainet_q6k_matmul_rm` (#1189) against the feed-order kernels they mirror. + * + * The same logical weight is laid out both ways — row-major `(o * blocksPerRow + b)` and + * feed-order `(b * outputDim + o)` — and both kernels must produce **bit-identical** outputs: + * per output row the accumulation order over blocks is the same, so any drift means the row + * addressing (the only thing that differs) is wrong. + */ +class RowMajorMatmulParityTest { + + private companion object { + const val BLOCK = 256 + const val Q4K_BPB = 144 + const val Q6K_BPB = 210 + } + + private fun bindRm(symbol: String): MethodHandle? { + val lookup = NativeLibraryLoader.lookup() ?: return null + val sym = lookup.find(symbol).orElse(null) ?: return null + val descriptor = FunctionDescriptor.ofVoid( + ValueLayout.ADDRESS, ValueLayout.JAVA_INT, // input, input_offset + ValueLayout.ADDRESS, ValueLayout.JAVA_INT, // weight, weight_byte_offset + ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, // input_dim, output_dim + ValueLayout.ADDRESS, ValueLayout.JAVA_INT, // output, output_offset + ) + return runCatching { Linker.nativeLinker().downcallHandle(sym, descriptor) }.getOrNull() + } + + private val q4kRm: MethodHandle? by lazy { bindRm("skainet_q4k_matmul_rm") } + private val q6kRm: MethodHandle? by lazy { bindRm("skainet_q6k_matmul_rm") } + + @BeforeTest + fun checkAvailable() { + assertTrue(NativeQ4KMatmulKernel.isAvailable(), "feed-order Q4_K kernel must be available") + assertTrue(NativeQ6KMatmulKernel.isAvailable(), "feed-order Q6_K kernel must be available") + assertNotNull(q4kRm, "skainet_q4k_matmul_rm must be bindable") + assertNotNull(q6kRm, "skainet_q6k_matmul_rm must be bindable") + } + + /** Random block bytes with the FP16 scale fields pinned to 1.0 so magnitudes stay sane. */ + private fun randomBlocks(numBlocks: Int, bytesPerBlock: Int, fp16At: IntArray, seed: Int): ByteArray { + val rng = Random(seed) + val bytes = ByteArray(numBlocks * bytesPerBlock) + rng.nextBytes(bytes) + for (block in 0 until numBlocks) { + val base = block * bytesPerBlock + for (off in fp16At) { + bytes[base + off] = 0x00 + bytes[base + off + 1] = 0x3C // 1.0f as binary16 + } + } + return bytes + } + + /** Permute row-major blocks `(o * bpr + b)` into feed order `(b * n + o)`. */ + private fun toFeedOrder(rowMajor: ByteArray, n: Int, bpr: Int, bpb: Int): ByteArray { + val out = ByteArray(rowMajor.size) + for (o in 0 until n) { + for (b in 0 until bpr) { + rowMajor.copyInto( + out, + destinationOffset = (b * n + o) * bpb, + startIndex = (o * bpr + b) * bpb, + endIndex = (o * bpr + b + 1) * bpb, + ) + } + } + return out + } + + private fun callRm( + handle: MethodHandle, + input: FloatArray, + rowMajorWeight: ByteArray, + weightByteOffset: Int, + inputDim: Int, + outputDim: Int, + ): FloatArray { + val out = FloatArray(outputDim) + Arena.ofConfined().use { arena -> + val inSeg = arena.allocate(inputDim.toLong() * 4, 4) + MemorySegment.copy(input, 0, inSeg, ValueLayout.JAVA_FLOAT, 0L, inputDim) + val wSeg = arena.allocate(rowMajorWeight.size.toLong(), 1) + MemorySegment.copy(rowMajorWeight, 0, wSeg, ValueLayout.JAVA_BYTE, 0L, rowMajorWeight.size) + val outSeg = arena.allocate(outputDim.toLong() * 4, 4) + handle.invoke(inSeg, 0, wSeg, weightByteOffset, inputDim, outputDim, outSeg, 0) + MemorySegment.copy(outSeg, ValueLayout.JAVA_FLOAT, 0L, out, 0, outputDim) + } + return out + } + + private fun assertQ4kParity(inputDim: Int, outputDim: Int, seed: Int, pad: Int = 0) { + val bpr = inputDim / BLOCK + val rowMajor = randomBlocks(bpr * outputDim, Q4K_BPB, intArrayOf(0, 2), seed) + val feed = toFeedOrder(rowMajor, outputDim, bpr, Q4K_BPB) + val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f } + + val expected = FloatArray(outputDim) + NativeQ4KMatmulKernel.matmul(input, 0, feed, 0, inputDim, outputDim, expected, 0) + + val padded = if (pad == 0) rowMajor else ByteArray(pad) + rowMajor + val got = callRm(q4kRm!!, input, padded, pad, inputDim, outputDim) + for (o in 0 until outputDim) { + assertEquals(expected[o].toRawBits(), got[o].toRawBits(), "Q4_K row $o diverged: ${expected[o]} vs ${got[o]}") + } + } + + private fun assertQ6kParity(inputDim: Int, outputDim: Int, seed: Int, pad: Int = 0) { + val bpr = inputDim / BLOCK + // Q6_K: d is the trailing FP16 at byte 208. + val rowMajor = randomBlocks(bpr * outputDim, Q6K_BPB, intArrayOf(208), seed) + val feed = toFeedOrder(rowMajor, outputDim, bpr, Q6K_BPB) + val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f } + + val expected = FloatArray(outputDim) + NativeQ6KMatmulKernel.matmul(input, 0, feed, 0, inputDim, outputDim, expected, 0) + + val padded = if (pad == 0) rowMajor else ByteArray(pad) + rowMajor + val got = callRm(q6kRm!!, input, padded, pad, inputDim, outputDim) + for (o in 0 until outputDim) { + assertEquals(expected[o].toRawBits(), got[o].toRawBits(), "Q6_K row $o diverged: ${expected[o]} vs ${got[o]}") + } + } + + @Test fun q4k_single_block_single_row() = assertQ4kParity(256, 1, seed = 42) + @Test fun q4k_single_block_multi_row() = assertQ4kParity(256, 16, seed = 7) + @Test fun q4k_multi_block_multi_row() = assertQ4kParity(1024, 64, seed = 123) + @Test fun q4k_llm_typical_shape() = assertQ4kParity(1536, 48, seed = 999) + @Test fun q4k_honors_weight_byte_offset() = assertQ4kParity(512, 8, seed = 17, pad = 257) + + @Test fun q6k_single_block_single_row() = assertQ6kParity(256, 1, seed = 41) + @Test fun q6k_single_block_multi_row() = assertQ6kParity(256, 16, seed = 8) + @Test fun q6k_multi_block_multi_row() = assertQ6kParity(1024, 64, seed = 321) + @Test fun q6k_honors_weight_byte_offset() = assertQ6kParity(512, 8, seed = 18, pad = 129) + + // ---- #1195: outputDim >= 512 engages the row-partition threading. Two oracles: ---- + + /** Threaded full-matrix call vs the feed-order kernel on permuted bytes (both threaded). */ + @Test fun q4k_threaded_parity_vs_feed_order() = assertQ4kParity(512, 1536, seed = 77) + @Test fun q6k_threaded_parity_vs_feed_order() = assertQ6kParity(512, 1536, seed = 78) + + /** + * Threaded full-matrix call vs 1536 independent single-row calls (each under the + * threshold, so single-threaded) — pins the partition arithmetic itself: every row of a + * threaded call must be bit-identical to that row computed alone. + */ + @Test + fun q4k_threaded_equals_per_row_calls() { + val inputDim = 512 + val n = 1536 + val bpr = inputDim / BLOCK + val rowMajor = randomBlocks(bpr * n, Q4K_BPB, intArrayOf(0, 2), seed = 91) + val input = FloatArray(inputDim) { Random(91 + it).nextFloat() - 0.5f } + + val full = callRm(q4kRm!!, input, rowMajor, 0, inputDim, n) + for (o in 0 until n step 97) { + val single = callRm(q4kRm!!, input, rowMajor, o * bpr * Q4K_BPB, inputDim, 1) + assertEquals(single[0].toRawBits(), full[o].toRawBits(), "Q4_K row $o: threaded diverged from solo") + } + } + + @Test + fun q6k_threaded_equals_per_row_calls() { + val inputDim = 512 + val n = 1536 + val bpr = inputDim / BLOCK + val rowMajor = randomBlocks(bpr * n, Q6K_BPB, intArrayOf(208), seed = 92) + val input = FloatArray(inputDim) { Random(92 + it).nextFloat() - 0.5f } + + val full = callRm(q6kRm!!, input, rowMajor, 0, inputDim, n) + for (o in 0 until n step 97) { + val single = callRm(q6kRm!!, input, rowMajor, o * bpr * Q6K_BPB, inputDim, 1) + assertEquals(single[0].toRawBits(), full[o].toRawBits(), "Q6_K row $o: threaded diverged from solo") + } + } +} diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/MappedFile.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/MappedFile.kt index c851fd516..b9e1df7c0 100644 --- a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/MappedFile.kt +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/MappedFile.kt @@ -26,6 +26,18 @@ public interface MappedFile : AutoCloseable { /** [length] bytes copied out of the mapping at [byteOffset], for kernels that need an array. */ public fun bytes(byteOffset: Long, length: Int): ByteArray + + /** + * A packed quantized tensor of [shape]/[encoding] viewing the mapping at [byteOffset] — zero + * heap bytes, blocks in canonical row-major file order (#1189). Returns `null` when this + * platform (or this encoding) has no off-heap packed representation; callers fall back to + * heap staging, which is the pre-#1189 behaviour for every packed tensor. + */ + public fun packedTensor( + byteOffset: Long, + shape: Shape, + encoding: sk.ainet.lang.tensor.storage.TensorEncoding, + ): TensorData<*, *>? = null } /** diff --git a/skainet-io/skainet-io-core/src/jvmAndroidMain/kotlin/sk/ainet/io/JvmMappedFile.kt b/skainet-io/skainet-io-core/src/jvmAndroidMain/kotlin/sk/ainet/io/JvmMappedFile.kt index d59a8d33e..39ba010b4 100644 --- a/skainet-io/skainet-io-core/src/jvmAndroidMain/kotlin/sk/ainet/io/JvmMappedFile.kt +++ b/skainet-io/skainet-io-core/src/jvmAndroidMain/kotlin/sk/ainet/io/JvmMappedFile.kt @@ -32,6 +32,32 @@ public class JvmMappedFile private constructor( return out } + /** + * Q4_K/Q6_K payloads as [sk.ainet.lang.tensor.data.BufferPackedTensorData] borrowing a slice + * of the one file mapping (#1189) — the packed counterpart of [denseFloats]: zero heap bytes, + * blocks left in canonical row-major file order for the buffer-reading kernels. + */ + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + override fun packedTensor( + byteOffset: Long, + shape: Shape, + encoding: sk.ainet.lang.tensor.storage.TensorEncoding, + ): TensorData<*, *>? = when (encoding) { + sk.ainet.lang.tensor.storage.TensorEncoding.Q4_K, + sk.ainet.lang.tensor.storage.TensorEncoding.Q6_K, + -> { + val length = checkNotNull(encoding.physicalBytes(shape.volume.toLong())) { + "${encoding.name} has size-determinate blocks" + } + sk.ainet.lang.tensor.data.BufferPackedTensorData( + shape, + sk.ainet.lang.memory.DirectBufferStorage.borrow(mmap.byteBufferAt(byteOffset, length)), + encoding, + ) + } + else -> null + } + /** Releases the channel; views already handed out keep working (the mapping outlives it). */ override fun close() { try { mmap.close() } finally { raf.close() } diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt index 31cb83389..f9357bbc8 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt @@ -29,6 +29,13 @@ public fun StreamingGGUFReader.planInput( prefillChunk: Int = PlanInput.DEFAULT_PREFILL_CHUNK, kvMode: KvCacheMode = KvCacheMode.BF16, nameMap: NameMap? = nameMap(), + /** + * The [sk.ainet.lang.memory.plan.WeightForm] the load will use, per tensor name — the same + * knob `StreamingGgufParametersLoader` takes, so the plan prices what the load will actually + * do (#1189: under `MAPPED` the servable encodings are budgeted against the page cache, not + * the heap). `null` for every tensor plans the pre-form default: everything heap-charged. + */ + formFor: (String) -> sk.ainet.lang.memory.plan.WeightForm? = { null }, ): PlanInput { val arch = fields["general.architecture"] as? String ?: "unknown" val name = fields["general.name"] as? String ?: arch @@ -41,6 +48,7 @@ public fun StreamingGGUFReader.planInput( format = format, elementCount = t.nElements, bytes = format.physicalBytes(t.nElements) ?: t.nBytes, + form = formFor(t.name), ) } val ctxUsed = ctx ?: geometry?.trainedContextLength ?: 2048 diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index 9c7b5dfb9..f6641edff 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -291,6 +291,39 @@ public class StreamingGgufParametersLoader( onProgress(current, total, tensorInfo.name) continue } + // #1189: a Q4_K/Q6_K tensor under MAPPED staging is the packed counterpart of the + // dense case above — served as a view over the file-backed pages, blocks left in + // canonical row-major order, never a heap ByteArray. Only when nothing asks for a + // different form: a dequantize/planes request or KERNEL_FEED order needs the bytes + // (or values) materialized anyway, and a platform without off-heap packed storage + // returns null and falls through to heap staging as before. + val mappedPacked: Tensor? = + if (mapped != null && tensorForm.residency == WeightResidency.MAPPED && + dtype == FP32::class && + tensorForm.order != WeightByteOrder.KERNEL_FEED && + tensorForm.encoding !is EncodingRequest.DequantizeTo && + !planesRequested(tensorForm) + ) { + val enc = when (tensorInfo.tensorType) { + GGMLQuantizationType.Q4_K -> sk.ainet.lang.tensor.storage.TensorEncoding.Q4_K + GGMLQuantizationType.Q6_K -> sk.ainet.lang.tensor.storage.TensorEncoding.Q6_K + else -> null + } + enc?.let { encoding -> + mapped.packedTensor(tensorInfo.absoluteDataOffset, shape, encoding)?.let { + @Suppress("UNCHECKED_CAST") + ctx.fromData(it as sk.ainet.lang.tensor.data.TensorData, dtype) + } + } + } else { + null + } + if (mappedPacked != null) { + onTensorLoaded(tensorInfo.name, mappedPacked) + current += 1 + onProgress(current, total, tensorInfo.name) + continue + } val rawBytes = mapped?.bytes(tensorInfo.absoluteDataOffset, tensorInfo.nBytes.toInt()) ?: reader.loadTensorData(tensorInfo) diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/MappedPackedStagingTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/MappedPackedStagingTest.kt new file mode 100644 index 000000000..c5b733d28 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/MappedPackedStagingTest.kt @@ -0,0 +1,95 @@ +package sk.ainet.io.gguf + +import java.io.File +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightResidency +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.BufferPackedTensorData +import sk.ainet.lang.tensor.data.Q4_KBlockTensorData +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 + +/** + * #1189: under `WeightResidency.MAPPED`, Q4_K/Q6_K payloads must be served as + * [BufferPackedTensorData] — a row-major view over the file mapping, zero heap bytes — while + * every other type keeps its pre-#1189 staging (heap packed data / mapped dense F32). Values + * must be identical to the heap load; StagingPolicyParityTest asserts that across policies, + * this test pins the *representation*. + */ +@OptIn(ExperimentalMemoryApi::class) +class MappedPackedStagingTest { + + private fun file(): File = SyntheticGguf.write( + SyntheticGguf.tensor("w_f32", GGMLQuantizationType.F32, elements = 1024), + SyntheticGguf.tensor("w_q4k", GGMLQuantizationType.Q4_K, elements = 1024), + SyntheticGguf.tensor("w_q6k", GGMLQuantizationType.Q6_K, elements = 1024), + SyntheticGguf.tensor("w_q80", GGMLQuantizationType.Q8_0, elements = 1024), + ) + + private fun load(f: File, form: WeightForm): Map> { + val ctx = DefaultDataExecutionContext() + val loaded = LinkedHashMap>() + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(f) }, + weightForm = form, + ).load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + @Test + fun `mapped staging keeps k-quant payloads off the heap, row-major`() { + val f = file() + try { + val mapped = load(f, WeightForm(residency = WeightResidency.MAPPED)) + + val q4k = mapped.getValue("w_q4k").data + assertTrue(q4k is BufferPackedTensorData, "Q4_K under MAPPED must be buffer-packed, got ${q4k::class.simpleName}") + assertEquals(TensorEncoding.Q4_K, q4k.encoding) + assertEquals(BlockOrder.ROW_MAJOR, q4k.blockOrder) + + val q6k = mapped.getValue("w_q6k").data + assertTrue(q6k is BufferPackedTensorData, "Q6_K under MAPPED must be buffer-packed, got ${q6k::class.simpleName}") + + // Unsupported-by-#1189 types keep their heap staging. + val q80 = mapped.getValue("w_q80").data + assertTrue(q80 !is BufferPackedTensorData, "Q8_0 has no buffer kernel yet; stays heap-staged") + + // Values equal the heap load, element for element. + val heap = load(f, WeightForm(residency = WeightResidency.HEAP)) + for (name in listOf("w_q4k", "w_q6k")) { + assertContentEquals( + heap.getValue(name).data.copyToFloatArray(), + mapped.getValue(name).data.copyToFloatArray(), + "values of $name", + ) + } + } finally { + f.delete() + } + } + + @Test + fun `heap residency is untouched`() { + val f = file() + try { + val heap = load(f, WeightForm(residency = WeightResidency.HEAP)) + val q4k = heap.getValue("w_q4k").data + assertTrue(q4k is Q4_KBlockTensorData, "HEAP residency keeps the heap tensor data") + assertTrue((q4k as PackedBlockStorage).packedData.isNotEmpty()) + } finally { + f.delete() + } + } +} diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StagingPolicyParityTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StagingPolicyParityTest.kt index abbfcaf08..6e44eae5d 100644 --- a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StagingPolicyParityTest.kt +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StagingPolicyParityTest.kt @@ -88,11 +88,18 @@ class StagingPolicyParityTest { mapped.getValue("w_f32").data is MmapFloatTensorData<*>, "mapped staging must not copy F32 onto the heap, got ${mapped.getValue("w_f32").data::class.simpleName}", ) - // packed tensors still arrive as packed block data: their kernels take arrays until #973 + // #1189: Q4_K/Q6_K under MAPPED are buffer-packed views over the mapping (zero heap + // bytes); other packed types still arrive as heap block data until they grow a + // buffer kernel. MappedPackedStagingTest pins the details. assertEquals( - heap.getValue("w_q4k").data::class.simpleName, + "BufferPackedTensorData", mapped.getValue("w_q4k").data::class.simpleName, - "packed staging is unchanged by the mapping", + "Q4_K under MAPPED stays off-heap (#1189)", + ) + assertEquals( + heap.getValue("w_q80").data::class.simpleName, + mapped.getValue("w_q80").data::class.simpleName, + "Q8_0 staging is unchanged by the mapping", ) } finally { f.delete() 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 9ac495315..bb110b1b6 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 @@ -1416,6 +1416,8 @@ public final class sk/ainet/lang/memory/plan/AllocationResolver { public static synthetic fun resolve$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; public final fun resolveTransient (Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;Lsk/ainet/lang/memory/ScopeKind;)Lsk/ainet/lang/memory/AllocationSpec; public static synthetic fun resolveTransient$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;Lsk/ainet/lang/memory/ScopeKind;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; + public final fun servesFromMapping (Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Z + public static synthetic fun servesFromMapping$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Z } public final class sk/ainet/lang/memory/plan/Budget { @@ -1583,8 +1585,8 @@ public final class sk/ainet/lang/memory/plan/KvCacheMode : java/lang/Enum { } public final class sk/ainet/lang/memory/plan/MemoryPlan { - public fun (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;J)V - public synthetic fun (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;JJ)V + public synthetic fun (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;JJILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Lsk/ainet/lang/memory/plan/PlanInput; public final fun component2 ()J public final fun component3 ()J @@ -1593,10 +1595,12 @@ public final class sk/ainet/lang/memory/plan/MemoryPlan { public final fun component6 ()J public final fun component7 ()Lsk/ainet/lang/memory/plan/Budget; public final fun component8 ()J - public final fun copy (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;J)Lsk/ainet/lang/memory/plan/MemoryPlan; - public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;JILjava/lang/Object;)Lsk/ainet/lang/memory/plan/MemoryPlan; + public final fun component9 ()J + public final fun copy (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;JJ)Lsk/ainet/lang/memory/plan/MemoryPlan; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;JJILjava/lang/Object;)Lsk/ainet/lang/memory/plan/MemoryPlan; public fun equals (Ljava/lang/Object;)Z public final fun getBudget ()Lsk/ainet/lang/memory/plan/Budget; + public final fun getBudgetedBytes ()J public final fun getFits ()Ljava/lang/Boolean; public final fun getFormConversionBytes ()J public final fun getForwardBytes ()J @@ -1609,6 +1613,8 @@ public final class sk/ainet/lang/memory/plan/MemoryPlan { public final fun getTotalBytes ()J public final fun getWeightsAsStoredBytes ()J public final fun getWeightsBytes ()J + public final fun getWeightsHeapBytes ()J + public final fun getWeightsMappedBytes ()J public fun hashCode ()I public final fun render ()Ljava/lang/String; public final fun suggestions ()Ljava/util/List; @@ -1625,8 +1631,8 @@ public final class sk/ainet/lang/memory/plan/MemoryPlans { public final fun formatBytes (J)Ljava/lang/String; public final fun forwardBytes (Lsk/ainet/lang/memory/plan/ModelGeometry;II)J public final fun kvElements (Lsk/ainet/lang/memory/plan/ModelGeometry;I)J - public final fun plan (Lsk/ainet/lang/memory/plan/PlanInput;Lsk/ainet/lang/memory/plan/Budget;)Lsk/ainet/lang/memory/plan/MemoryPlan; - public static synthetic fun plan$default (Lsk/ainet/lang/memory/plan/MemoryPlans;Lsk/ainet/lang/memory/plan/PlanInput;Lsk/ainet/lang/memory/plan/Budget;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/MemoryPlan; + public final fun plan (Lsk/ainet/lang/memory/plan/PlanInput;Lsk/ainet/lang/memory/plan/Budget;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Lsk/ainet/lang/memory/plan/MemoryPlan; + public static synthetic fun plan$default (Lsk/ainet/lang/memory/plan/MemoryPlans;Lsk/ainet/lang/memory/plan/PlanInput;Lsk/ainet/lang/memory/plan/Budget;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/MemoryPlan; } public final class sk/ainet/lang/memory/plan/ModelGeometry { @@ -1688,16 +1694,19 @@ public final class sk/ainet/lang/memory/plan/PlanInput$Companion { } public final class sk/ainet/lang/memory/plan/PlanLine { - public fun (Ljava/lang/String;Ljava/lang/String;JZ)V + public fun (Ljava/lang/String;Ljava/lang/String;JZZ)V + public synthetic fun (Ljava/lang/String;Ljava/lang/String;JZZILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; public final fun component2 ()Ljava/lang/String; public final fun component3 ()J public final fun component4 ()Z - public final fun copy (Ljava/lang/String;Ljava/lang/String;JZ)Lsk/ainet/lang/memory/plan/PlanLine; - public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanLine;Ljava/lang/String;Ljava/lang/String;JZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanLine; + public final fun component5 ()Z + public final fun copy (Ljava/lang/String;Ljava/lang/String;JZZ)Lsk/ainet/lang/memory/plan/PlanLine; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanLine;Ljava/lang/String;Ljava/lang/String;JZZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanLine; public fun equals (Ljava/lang/Object;)Z public final fun getBytes ()J public final fun getDetail ()Ljava/lang/String; + public final fun getMapped ()Z public final fun getResident ()Z public final fun getSection ()Ljava/lang/String; public fun hashCode ()I @@ -1857,13 +1866,15 @@ public final class sk/ainet/lang/memory/plan/ProfiledPlan { public final class sk/ainet/lang/memory/plan/StorageCapabilities { public static final field Companion Lsk/ainet/lang/memory/plan/StorageCapabilities$Companion; - public fun (ZZ)V - public synthetic fun (ZZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (ZZLjava/util/Set;)V + public synthetic fun (ZZLjava/util/Set;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Z public final fun component2 ()Z - public final fun copy (ZZ)Lsk/ainet/lang/memory/plan/StorageCapabilities; - public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/StorageCapabilities;ZZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/StorageCapabilities; + public final fun component3 ()Ljava/util/Set; + public final fun copy (ZZLjava/util/Set;)Lsk/ainet/lang/memory/plan/StorageCapabilities; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/StorageCapabilities;ZZLjava/util/Set;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/StorageCapabilities; public fun equals (Ljava/lang/Object;)Z + public final fun getMappedServableEncodings ()Ljava/util/Set; public final fun getSupportsMappedFiles ()Z public final fun getSupportsOffHeap ()Z public fun hashCode ()I @@ -1874,6 +1885,7 @@ public final class sk/ainet/lang/memory/plan/StorageCapabilities$Companion { public final fun current ()Lsk/ainet/lang/memory/plan/StorageCapabilities; public final fun getFULL ()Lsk/ainet/lang/memory/plan/StorageCapabilities; public final fun getHEAP_ONLY ()Lsk/ainet/lang/memory/plan/StorageCapabilities; + public final fun getMAPPED_SERVABLE_DEFAULT ()Ljava/util/Set; } public final class sk/ainet/lang/memory/plan/Suggestion { @@ -4729,6 +4741,30 @@ public final class sk/ainet/lang/tensor/data/BitNetPlanesTensorData$Companion { public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/BitNetPlanesTensorData; } +public final class sk/ainet/lang/tensor/data/BufferPackedTensorData : sk/ainet/lang/tensor/data/TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { + public fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/storage/TensorEncoding;)V + public fun copyToFloatArray ()[F + public fun dequantizeBlock (I[FI)V + public fun get ([I)Ljava/lang/Float; + public synthetic fun get ([I)Ljava/lang/Object; + public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; + public fun getBlockSize ()I + public fun getElementCount ()J + public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; + public fun getPackedData ()[B + public fun getPackedView ()Lsk/ainet/lang/memory/TensorView; + public fun getPhysicalBytes ()J + public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public final fun getStorage ()Lsk/ainet/lang/memory/Storage; + public fun getView ()Lsk/ainet/lang/memory/TensorView; + public fun set ([IF)V + public synthetic fun set ([ILjava/lang/Object;)V + public fun toFloatArray ()[F + public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; +} + public final class sk/ainet/lang/tensor/data/DenseFloatArrayTensorData : sk/ainet/lang/tensor/data/FloatArrayTensorData { public fun (Lsk/ainet/lang/tensor/Shape;[F)V public fun copyToFloatArray ()[F @@ -4924,6 +4960,7 @@ public final class sk/ainet/lang/tensor/data/MmapFloatTensorData : sk/ainet/lang public final class sk/ainet/lang/tensor/data/MmapTensorSource : java/lang/AutoCloseable { public static final field Companion Lsk/ainet/lang/tensor/data/MmapTensorSource$Companion; public fun (Ljava/nio/MappedByteBuffer;)V + public final fun byteBufferAt (JJ)Ljava/nio/ByteBuffer; public fun close ()V public final fun floatTensorAt (JLsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/data/MmapFloatTensorData; public final fun isLoaded ()Z diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt index 3f28fbdf6..f6393f377 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt @@ -6,6 +6,7 @@ import sk.ainet.lang.memory.Format import sk.ainet.lang.memory.PlatformStorage import sk.ainet.lang.memory.ScopeKind import sk.ainet.lang.tensor.storage.MemoryDomain +import sk.ainet.lang.tensor.storage.TensorEncoding /** * What the running platform's storage can actually do — the third input of [AllocationResolver], @@ -15,8 +16,21 @@ import sk.ainet.lang.tensor.storage.MemoryDomain public data class StorageCapabilities( val supportsMappedFiles: Boolean, val supportsOffHeap: Boolean = true, + /** + * Encodings whose bytes the runtime can actually *serve* from a file mapping — a tensor is + * only truly mapped when a loader emits a file-backed representation for it AND a kernel can + * read it there. Today (#921 dense F32, #1189 buffer-packed Q4_K/Q6_K) that is exactly + * [MAPPED_SERVABLE_DEFAULT]; every other encoding under `WeightResidency.MAPPED` falls back + * to heap staging, and a plan that assumed otherwise would under-count the heap it is about + * to fill (the mirror image of the #1116 dequantize-surprise). + */ + val mappedServableEncodings: Set = MAPPED_SERVABLE_DEFAULT, ) { public companion object { + /** What the loaders can serve from mapped pages today: dense FP32, Q4_K, Q6_K (#1189). */ + public val MAPPED_SERVABLE_DEFAULT: Set = + setOf(TensorEncoding.Dense(4), TensorEncoding.Q4_K, TensorEncoding.Q6_K) + /** The platform this code is running on. */ public fun current(): StorageCapabilities = StorageCapabilities( supportsMappedFiles = PlatformStorage.supportsMappedFiles, @@ -53,16 +67,33 @@ public object AllocationResolver { * the file it no longer matches. Everything else falls to [PlannerProfile.domainFor] over the * bytes actually held, so a dequantized giant goes off-heap and a small bias stays on it. */ - public fun resolve( + /** + * Whether [weight] will really be served from file-backed pages: the form asks for + * [WeightResidency.MAPPED], the platform can map, the bytes are the file's bytes (no + * re-encode, no re-order), **and** the encoding is one the runtime can serve from a mapping + * ([StorageCapabilities.mappedServableEncodings] — dense F32 since #921, Q4_K/Q6_K since + * #1189). This is the predicate the plan uses to budget a weight against the page cache + * instead of the heap, so it must not overclaim. + */ + public fun servesFromMapping( weight: PlanTensor, - profile: PlannerProfile, platform: StorageCapabilities = StorageCapabilities.current(), - ): AllocationSpec { + ): Boolean { val form = weight.form val fileBytesAreTheBytes = form == null || (form.encoding == EncodingRequest.KeepAsStored && form.order == WeightByteOrder.AS_STORED) - val wantsMapped = form?.residency == WeightResidency.MAPPED - val mapped = wantsMapped && platform.supportsMappedFiles && fileBytesAreTheBytes + return form?.residency == WeightResidency.MAPPED && + platform.supportsMappedFiles && + fileBytesAreTheBytes && + weight.format.encoding in platform.mappedServableEncodings + } + + public fun resolve( + weight: PlanTensor, + profile: PlannerProfile, + platform: StorageCapabilities = StorageCapabilities.current(), + ): AllocationSpec { + val mapped = servesFromMapping(weight, platform) val domain = if (mapped) MemoryDomain.MMAP_FILE else fallbackDomain(weight.residentBytes, profile, platform) return AllocationSpec( format = residentFormat(weight), @@ -116,6 +147,10 @@ public object AllocationResolver { "form asks MAPPED but the weight is re-encoded at load — a copy cannot be paged from the file" form?.residency == WeightResidency.MAPPED && form.order == WeightByteOrder.KERNEL_FEED -> "form asks MAPPED but kernel-feed order is a load-time copy" + form?.residency == WeightResidency.MAPPED && + weight.format.encoding !in platform.mappedServableEncodings -> + "form asks MAPPED but no loader/kernel serves ${weight.format.encoding.name} " + + "from a mapping yet (#1189 covers dense F32, Q4_K, Q6_K) — heap staging" else -> "resident ${MemoryPlans.formatBytes(weight.residentBytes)} vs off-heap threshold " + MemoryPlans.formatBytes(profile.offHeapThresholdBytes) + 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 d99b80d37..660ee5c50 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 @@ -142,9 +142,19 @@ public data class Budget(val bytes: Long, val description: String) { } } -/** One line of the plan: what, how much, and whether it is resident for the whole session. */ +/** + * One line of the plan: what, how much, and whether it is resident for the whole session. + * A [mapped] line is file-backed page cache — resident in RSS but evictable, and **not** charged + * against the heap [Budget] (#1189: a mapped 1 GB model runs under a 256 MB ART cap). + */ @ExperimentalMemoryApi -public data class PlanLine(val section: String, val detail: String, val bytes: Long, val resident: Boolean) +public data class PlanLine( + val section: String, + val detail: String, + val bytes: Long, + val resident: Boolean, + val mapped: Boolean = false, +) /** A concrete way to make a plan fit, with the bytes it saves. */ @ExperimentalMemoryApi @@ -172,36 +182,71 @@ public data class MemoryPlan( * forms existed — hence the default. */ val weightsAsStoredBytes: Long = weightsBytes, + /** + * The part of [weightsBytes] served from file-backed pages (#1189): resident in RSS but + * evictable and outside the managed heap, so it is charged against device RAM/page cache + * rather than the [budget]. Zero for a fully heap-staged model — every plan before #1189. + */ + val weightsMappedBytes: Long = 0L, ) { /** * Bytes the resolved forms add to the weights — a dequantization's price, made visible before * it is paid rather than discovered as an OOM. Zero when the weights are held as stored. */ val formConversionBytes: Long get() = weightsBytes - weightsAsStoredBytes + + /** The part of [weightsBytes] that really lands on the managed heap. */ + val weightsHeapBytes: Long get() = weightsBytes - weightsMappedBytes + + /** The full footprint — heap sections plus mapped pages. What RSS converges to, not what the heap holds. */ val totalBytes: Long get() = weightsBytes + kvBytes + forwardBytes + headroomBytes + + /** + * What is charged against the [budget]: everything except the mapped weight pages, which the + * OS pages in and evicts against device RAM (#1189 measured: 1.0 GB mapped, 566 KB heap, + * decode under a 256 MB cap). Equal to [totalBytes] when nothing is mapped. + */ + val budgetedBytes: Long get() = totalBytes - weightsMappedBytes val residentBytes: Long get() = weightsBytes + kvBytes - /** `true` when a budget is set and the total fits; `null` without a budget. */ - val fits: Boolean? get() = budget?.let { totalBytes <= it.bytes } + /** `true` when a budget is set and the budget-charged total fits; `null` without a budget. */ + val fits: Boolean? get() = budget?.let { budgetedBytes <= it.bytes } val lines: List - get() = listOf( - PlanLine( - "weights", - if (formConversionBytes == 0L) "Mapped, packed" - else "re-encoded at load (+${MemoryPlans.formatBytes(formConversionBytes)})", - weightsBytes, - resident = true, - ), - PlanLine("kv cache", input.kvMode.label + " @ ctx ${input.ctx}", kvBytes, resident = true), - PlanLine("forward", "prefill chunk ${input.prefillChunk}", forwardBytes, resident = false), - PlanLine("heap", "headroom", headroomBytes, resident = false), - ) + get() = buildList { + if (weightsMappedBytes > 0L) { + add(PlanLine("weights", "mapped, as stored", weightsMappedBytes, resident = true, mapped = true)) + if (weightsHeapBytes > 0L) { + add( + PlanLine( + "weights", + if (formConversionBytes == 0L) "heap-staged, packed" + else "re-encoded at load (+${MemoryPlans.formatBytes(formConversionBytes)})", + weightsHeapBytes, + resident = true, + ), + ) + } + } else { + add( + PlanLine( + "weights", + if (formConversionBytes == 0L) "heap-staged, packed" + else "re-encoded at load (+${MemoryPlans.formatBytes(formConversionBytes)})", + weightsBytes, + resident = true, + ), + ) + } + add(PlanLine("kv cache", input.kvMode.label + " @ ctx ${input.ctx}", kvBytes, resident = true)) + add(PlanLine("forward", "prefill chunk ${input.prefillChunk}", forwardBytes, resident = false)) + add(PlanLine("heap", "headroom", headroomBytes, resident = false)) + } /** At least two concrete suggestions with their savings when the plan does not fit (M0-F3). */ public fun suggestions(): List { val b = budget ?: return emptyList() - if (totalBytes <= b.bytes) return emptyList() + if (budgetedBytes <= b.bytes) return emptyList() val out = ArrayList() if (input.kvMode == KvCacheMode.BF16 && kvBytesAlternate < kvBytes) { out += Suggestion("--kv turboquant", kvBytes - kvBytesAlternate) @@ -209,7 +254,7 @@ public data class MemoryPlan( val halfCtx = (input.ctx / 2).coerceAtLeast(1) if (halfCtx < input.ctx) { val half = MemoryPlans.plan(input.copy(ctx = halfCtx), budget) - out += Suggestion("--ctx $halfCtx", totalBytes - half.totalBytes) + out += Suggestion("--ctx $halfCtx", budgetedBytes - half.budgetedBytes) } if (formConversionBytes > 0) { // Worth saying first: unlike ctx or KV mode, this cost was not asked for — it is what @@ -219,7 +264,7 @@ public data class MemoryPlan( formConversionBytes, ) } - val over = totalBytes - b.bytes + val over = budgetedBytes - b.bytes out += Suggestion("a smaller model: weights must shrink by ≥ ${MemoryPlans.formatBytes(over)} (e.g. a lower-bit quantization of the same model)", over) return out } @@ -232,15 +277,21 @@ public data class MemoryPlan( append(" · ctx "); append(input.ctx); append('\n') 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.mapped) append(" mapped (page cache, evictable — not heap)") + else if (l.resident) append(" resident") 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)) + val budgetLabel = if (weightsMappedBytes > 0L) "total heap" else "total" + append(" "); append(budgetLabel.padEnd(36)); append(MemoryPlans.formatBytes(budgetedBytes).padStart(10)) val b = budget if (b != null) { append(" of "); append(MemoryPlans.formatBytes(b.bytes)); append(if (fits == true) " ✔ fits" else " ✘ does not fit") append('\n') + if (weightsMappedBytes > 0L) { + append(" mapped weights (").append(MemoryPlans.formatBytes(weightsMappedBytes)) + append(") page against device RAM, not this budget (#1189)\n") + } val s = suggestions() if (s.isNotEmpty()) { append(" suggestions: "); append(s.joinToString(" · ") { "${it.text} (−${MemoryPlans.formatBytes(it.savesBytes)})" }); append('\n') @@ -262,22 +313,31 @@ public object MemoryPlans { /** * Build the plan. Estimates: - * - weights: sum of the packed byte sizes (they are touched every token, so counted resident); + * - weights: sum of the packed byte sizes (they are touched every token, so counted resident) — + * split by [AllocationResolver.servesFromMapping] into heap-charged bytes and mapped bytes, + * because a mapped weight pages against device RAM, not the budget (#1189); * - kv cache: `layers × 2 × ctx × kvHeads × (headDim + valueDim)/2 × mode bytes`; * - forward slab for a chunk of `T = min(prefillChunk, ctx)` tokens, FP32: * `T × (4·emb + 3·ffn + heads·ctx) × 4 B` (residual stream, attention projections, gated FFN * intermediates, attention scores over the context) plus one `vocab × 4 B` logits row; * - heap headroom: [HEAP_HEADROOM_BYTES]. */ - public fun plan(input: PlanInput, budget: Budget? = null): MemoryPlan { + public fun plan( + input: PlanInput, + budget: Budget? = null, + platform: StorageCapabilities = StorageCapabilities.current(), + ): MemoryPlan { val weights = input.weights.sumOf { it.residentBytes } + val weightsMapped = input.weights + .filter { AllocationResolver.servesFromMapping(it, platform) } + .sumOf { it.residentBytes } val weightsAsStored = input.weights.sumOf { it.bytes } 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.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, weightsAsStored) + return MemoryPlan(input, weights, kv, kvAlt, forward, HEAP_HEADROOM_BYTES, budget, weightsAsStored, weightsMapped) } public fun kvElements(g: ModelGeometry, ctx: Int): Long = diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt index bf75d8e4b..19fe997da 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt @@ -64,10 +64,12 @@ public data class PlannerProfile( val base = MemoryPlans.plan(input.copy(prefillChunk = prefillChunk, kvMode = kvMode), budget) val notes = ArrayList() var plan = base - val share = if (budget.bytes > 0) base.totalBytes.toDouble() / budget.bytes else Double.MAX_VALUE + // Budget pressure comes from what is charged against the budget — mapped weights page + // against device RAM, not the heap (#1189), so they must not trigger KV quantization. + val share = if (budget.bytes > 0) base.budgetedBytes.toDouble() / budget.bytes else Double.MAX_VALUE if (kvMode != KvCacheMode.TURBOQUANT_4 && share > kvAutoQuantizeAbove) { val quantized = MemoryPlans.plan(base.input.copy(kvMode = KvCacheMode.TURBOQUANT_4), budget) - if (quantized.totalBytes < base.totalBytes) { + if (quantized.budgetedBytes < base.budgetedBytes) { plan = quantized notes += "KV cache switched to ${KvCacheMode.TURBOQUANT_4.label}: the plan needed " + "${percent(share)} of the budget (over ${percent(kvAutoQuantizeAbove)}), saving " + diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MappedBudgetPlanTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MappedBudgetPlanTest.kt new file mode 100644 index 000000000..129a058c0 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MappedBudgetPlanTest.kt @@ -0,0 +1,111 @@ +package sk.ainet.lang.memory.plan + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertNotEquals +import kotlin.test.assertTrue +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.tensor.storage.MemoryDomain +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 + +/** + * #1189's planner half: weights served from file-backed pages are charged against device RAM, not + * the heap [Budget]. Pinned by the Pixel 8a measurement — a 1.0 GB mapped Q4_K_M model decoded + * under a 256 MB ART cap with 566 KB of weight heap; a plan that says "does not fit" about that + * run is wrong, and this test is what keeps it from saying so again. + */ +@OptIn(ExperimentalMemoryApi::class) +class MappedBudgetPlanTest { + + private val platform = StorageCapabilities.FULL + + private fun q4k(name: String, bytes: Long, form: WeightForm?) = PlanTensor( + name = name, + id = null, + format = Format(FP32, TensorEncoding.Q4_K), + elementCount = bytes / 144 * 256, + bytes = bytes, + form = form, + ) + + private fun input(weights: List) = PlanInput( + modelName = "m", architecture = "llama", weights = weights, geometry = null, ctx = 512, + ) + + @Test + fun mapped_weights_are_not_charged_against_the_heap_budget() { + val gig = 1024L * 1024 * 1024 + val cap = 256L * 1024 * 1024 + val mappedForm = WeightForm(residency = WeightResidency.MAPPED) + val plan = MemoryPlans.plan(input(listOf(q4k("w", gig, mappedForm))), Budget.of(cap), platform) + + assertEquals(gig, plan.weightsMappedBytes) + assertEquals(0L, plan.weightsHeapBytes) + assertEquals(plan.totalBytes - gig, plan.budgetedBytes) + assertEquals(true, plan.fits, "1 GB mapped weights under a 256 MB cap FIT — measured, #1189:\n" + plan.render()) + assertTrue(plan.suggestions().isEmpty(), "a fitting plan suggests nothing") + + val render = plan.render() + assertTrue("mapped" in render && "✔ fits" in render, render) + assertTrue("total heap" in render, render) + } + + @Test + fun heap_staged_weights_still_count_and_still_overflow() { + val gig = 1024L * 1024 * 1024 + val cap = 256L * 1024 * 1024 + val plan = MemoryPlans.plan(input(listOf(q4k("w", gig, form = null))), Budget.of(cap), platform) + + assertEquals(0L, plan.weightsMappedBytes) + assertEquals(plan.totalBytes, plan.budgetedBytes) + assertEquals(false, plan.fits) + assertFalse(plan.suggestions().isEmpty(), "an over-budget plan must suggest a way out") + } + + @Test + fun encodings_nobody_serves_from_a_mapping_stay_heap_charged_even_under_MAPPED() { + val bytes = 100L * 1024 * 1024 + val mappedForm = WeightForm(residency = WeightResidency.MAPPED) + val q80 = PlanTensor("w", null, Format(FP32, TensorEncoding.Q8_0), bytes / 34 * 32, bytes, mappedForm) + val plan = MemoryPlans.plan(input(listOf(q80)), Budget.of(256L * 1024 * 1024), platform) + + assertEquals(0L, plan.weightsMappedBytes, "Q8_0 has no mapped-servable representation (#1189 covers Q4_K/Q6_K)") + + // and the resolver agrees — the plan and the load must tell the same story + val spec = AllocationResolver.resolve(q80, PlannerProfile.MOBILE_2GB, platform) + assertNotEquals(MemoryDomain.MMAP_FILE, spec.domain) + assertTrue("no loader/kernel serves Q8_0" in AllocationResolver.explain(q80, PlannerProfile.MOBILE_2GB, platform)) + } + + @Test + fun a_mix_splits_into_a_mapped_and_a_heap_line() { + val mappedForm = WeightForm(residency = WeightResidency.MAPPED) + val big = q4k("big", 512L * 1024 * 1024, mappedForm) + val q80 = PlanTensor( + "small", null, Format(FP32, TensorEncoding.Q8_0), + (10L * 1024 * 1024) / 34 * 32, 10L * 1024 * 1024, mappedForm, + ) + val plan = MemoryPlans.plan(input(listOf(big, q80)), Budget.of(256L * 1024 * 1024), platform) + + assertEquals(512L * 1024 * 1024, plan.weightsMappedBytes) + assertEquals(10L * 1024 * 1024, plan.weightsHeapBytes) + val weightLines = plan.lines.filter { it.section == "weights" } + assertEquals(2, weightLines.size) + assertTrue(weightLines[0].mapped && !weightLines[1].mapped) + } + + @Test + fun without_mapping_support_everything_is_heap_charged() { + val mappedForm = WeightForm(residency = WeightResidency.MAPPED) + val plan = MemoryPlans.plan( + input(listOf(q4k("w", 512L * 1024 * 1024, mappedForm))), + Budget.of(256L * 1024 * 1024), + StorageCapabilities.HEAP_ONLY, + ) + assertEquals(0L, plan.weightsMappedBytes) + assertEquals(false, plan.fits) + } +} diff --git a/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/BufferPackedTensorData.kt b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/BufferPackedTensorData.kt new file mode 100644 index 000000000..df332a72a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/BufferPackedTensorData.kt @@ -0,0 +1,156 @@ +package sk.ainet.lang.tensor.data + +import java.nio.ByteBuffer +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.DirectBufferStorage +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.MappedBufferStorage +import sk.ainet.lang.memory.PackedBlockDecoder +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 + +/** + * A packed quantized weight whose bytes live **off the managed heap** — in a [MappedBufferStorage] + * (mmap'd file pages) or a [DirectBufferStorage] — in canonical GGUF row-major block order (#1189). + * + * This is the packed sibling of [MmapFloatTensorData]: where that class keeps dense F32 weights + * out of ART's heap cap, this one does it for Q4_K/Q6_K payloads, which under + * `WeightResidency.MAPPED` used to materialize as heap `ByteArray`s and were the reason a 1 GB + * Q4_K_M model died at load under a 256 MB cap while the mapped machinery itself held major + * faults at zero (#1130). + * + * The fast path is [packedView]: a `BLOCKED_ROW_MAJOR` view over the off-heap storage that the + * buffer-reading kernels (JNI direct-buffer, FFM) consume without copying a byte. Element access + * ([get], [dequantizeBlock], [toFloatArray]) goes through a single-block scratch delegate — + * correct but per-block-copy slow, and **not thread-safe**; it exists for embedding-lookup-style + * reads and references, not for matmuls. + * + * [get] mirrors the heap classes' raw-code semantics (`Q4_KBlockTensorData.get` returns the + * quantization *code*, not the value — PackedBlockStorage's documented source-compat quirk), so + * that the "staging never changes the numbers" invariant holds across HEAP↔MAPPED loads of the + * same file. Decoded values come from [packedView]`.get` (rule 4) or [dequantizeBlock]. + * + * [packedData] deliberately throws: handing out a heap `ByteArray` is exactly what this class + * exists to avoid. Readers that need bytes go through [packedView]'s storage; readers that need + * values go through [dequantizeBlock]. + */ +@ExperimentalMemoryApi +public class BufferPackedTensorData( + initialShape: Shape, + /** Off-heap storage holding exactly this tensor's packed blocks (payload only, offset 0). */ + public val storage: Storage, + override val encoding: TensorEncoding, +) : TensorData, PackedBlockStorage { + + override val shape: Shape = Shape(initialShape.dimensions.copyOf()) + private val strides: IntArray = shape.computeStrides() + + override val blockSize: Int + private val bytesPerBlock: Int + + /** Decodes one block out of [scratch]; refilled from the buffer before each decode. */ + private val scratch: ByteArray + private val scratchDecoder: PackedBlockStorage + private val scratchData: TensorData<*, *> + + private val buf: ByteBuffer = when (storage) { + is MappedBufferStorage -> storage.buffer() + is DirectBufferStorage -> storage.buffer() + else -> throw IllegalArgumentException( + "BufferPackedTensorData needs buffer-backed off-heap storage " + + "(MappedBufferStorage or DirectBufferStorage), got ${storage::class.simpleName}" + ) + } + + init { + when (encoding) { + TensorEncoding.Q4_K -> { + blockSize = TensorEncoding.Q4_K.BLOCK_SIZE + bytesPerBlock = TensorEncoding.Q4_K.BYTES_PER_BLOCK + scratch = ByteArray(bytesPerBlock) + val d = Q4_KBlockTensorData(Shape(blockSize), scratch) + scratchDecoder = d + scratchData = d + } + TensorEncoding.Q6_K -> { + blockSize = TensorEncoding.Q6_K.BLOCK_SIZE + bytesPerBlock = TensorEncoding.Q6_K.BYTES_PER_BLOCK + scratch = ByteArray(bytesPerBlock) + val d = Q6_KBlockTensorData(Shape(blockSize), scratch) + scratchDecoder = d + scratchData = d + } + else -> throw IllegalArgumentException( + "BufferPackedTensorData supports Q4_K and Q6_K (#1189); got ${encoding.name}. " + + "Other block formats keep their heap TensorData until a buffer kernel exists." + ) + } + require(shape.volume % blockSize == 0) { + "shape $shape (${shape.volume} elements) is not a whole number of $blockSize-element " + + "${encoding.name} blocks" + } + require(storage.sizeBytes == blockCount.toLong() * bytesPerBlock) { + "storage holds ${storage.sizeBytes} bytes; ${encoding.name} $shape needs exactly " + + "${blockCount.toLong() * bytesPerBlock} (payload only, no trailer)" + } + } + + override val blockCount: Int get() = shape.volume / blockSize + override val blockOrder: BlockOrder get() = BlockOrder.ROW_MAJOR + override val physicalBytes: Long get() = storage.sizeBytes + + override val packedData: ByteArray + get() = throw UnsupportedOperationException( + "BufferPackedTensorData keeps its ${encoding.name} bytes off-heap (#1189) — there is " + + "no heap ByteArray to hand out. Kernels take packedView's storage; element readers " + + "use dequantizeBlock/get." + ) + + override val packedView: TensorView + get() = TensorView.packed( + storage = storage, + shape = shape, + encoding = encoding, + decoder = PackedBlockDecoder(this), + blockOrder = blockOrder, + ) + + override val view: TensorView get() = packedView + + override fun dequantizeBlock(blockIdx: Int, output: FloatArray, outputOffset: Int) { + require(blockIdx in 0 until blockCount) { "Block index $blockIdx out of bounds (0..$blockCount)" } + val d = buf.duplicate() + d.position(blockIdx * bytesPerBlock) + d.get(scratch, 0, bytesPerBlock) + scratchDecoder.dequantizeBlock(0, output, outputOffset) + } + + override fun get(vararg indices: Int): Float { + require(indices.size == shape.dimensions.size) { + "Number of indices (${indices.size}) must match tensor dimensions (${shape.dimensions.size})" + } + var flat = 0 + for (i in indices.indices) { + val idx = indices[i] + require(idx >= 0 && idx < shape.dimensions[i]) { + "Index $idx out of bounds for dimension $i with size ${shape.dimensions[i]}" + } + flat += idx * strides[i] + } + val block = flat / blockSize + val within = flat % blockSize + val d = buf.duplicate() + d.position(block * bytesPerBlock) + d.get(scratch, 0, bytesPerBlock) + // Raw code, as the heap classes return it — see the class KDoc. Decoded values are + // packedView.get's job. + return (scratchData.get(within) as Number).toFloat() + } + + override fun set(vararg indices: Int, value: Float): Unit = + throw UnsupportedOperationException("BufferPackedTensorData is read-only (mapped/borrowed weight bytes)") +} diff --git a/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt index d3c609fa5..1f66fef7f 100644 --- a/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt +++ b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt @@ -135,6 +135,24 @@ public class MmapTensorSource( return MmapFloatTensorData(shape, slice) } + /** + * A direct [ByteBuffer] slice of the mapping — `[byteOffset, byteOffset + length)`, little + * endian, position 0. Zero-copy: reads hit the file-backed pages. This is what packed + * (quantized) tensors are served from under mapped staging (#1189); dense F32 tensors use + * [floatTensorAt]. + */ + public fun byteBufferAt(byteOffset: Long, length: Long): ByteBuffer { + require(byteOffset >= 0 && length >= 0) { "byteOffset and length must be non-negative" } + require(byteOffset + length <= mappedBuffer.capacity()) { + "Region [$byteOffset, ${byteOffset + length}) exceeds buffer capacity ${mappedBuffer.capacity()}" + } + // Not chained: Android's pre-Java-9 nio signatures return Buffer (see floatTensorAt). + val dup = mappedBuffer.duplicate() + dup.position(byteOffset.toInt()) + dup.limit((byteOffset + length).toInt()) + return dup.slice().order(ByteOrder.LITTLE_ENDIAN) + } + /** * Force the mapped memory to be loaded into physical memory. * This can improve first-access performance but uses more memory. diff --git a/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/tensor/data/BufferPackedTensorDataTest.kt b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/tensor/data/BufferPackedTensorDataTest.kt new file mode 100644 index 000000000..4e37c2cea --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/tensor/data/BufferPackedTensorDataTest.kt @@ -0,0 +1,110 @@ +package sk.ainet.lang.tensor.data + +import java.nio.file.Files +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.DirectBufferStorage +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.MappedBufferStorage +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * #1189: a [BufferPackedTensorData] over off-heap bytes must decode exactly like the heap + * `TensorData` for the same bytes — mapped file or direct buffer, Q4_K and Q6_K. + */ +@OptIn(ExperimentalMemoryApi::class) +class BufferPackedTensorDataTest { + + private fun randomPayload(numBlocks: Int, bytesPerBlock: Int, fp16At: IntArray, seed: Int): ByteArray { + val rng = Random(seed) + val bytes = ByteArray(numBlocks * bytesPerBlock) + rng.nextBytes(bytes) + for (block in 0 until numBlocks) { + for (off in fp16At) { + bytes[block * bytesPerBlock + off] = 0x00 + bytes[block * bytesPerBlock + off + 1] = 0x3C + } + } + return bytes + } + + private fun directStorage(bytes: ByteArray): DirectBufferStorage { + val s = DirectBufferStorage.allocate(bytes.size) + s.buffer().put(bytes) + return s + } + + @Test + fun q4k_direct_buffer_decodes_like_heap() { + val shape = Shape(4, 512) // 2 blocks per row, 8 blocks + val bytes = randomPayload(8, 144, intArrayOf(0, 2), seed = 5) + val heap = Q4_KBlockTensorData(shape, bytes) + val buf = BufferPackedTensorData(shape, directStorage(bytes), TensorEncoding.Q4_K) + + assertContentEquals(heap.toFloatArray(), buf.toFloatArray()) + // TensorData.get mirrors the heap class's raw-code semantics (see class KDoc) + assertEquals(heap.get(3, 511).toFloat(), buf.get(3, 511)) + assertContentEquals(heap.copyToFloatArray(), buf.copyToFloatArray()) + assertEquals(bytes.size.toLong(), buf.physicalBytes) + assertEquals(BlockOrder.ROW_MAJOR, buf.blockOrder) + } + + @Test + fun q6k_mapped_file_decodes_like_heap() { + val shape = Shape(2, 512) + val bytes = randomPayload(4, 210, intArrayOf(208), seed = 9) + val heap = Q6_KBlockTensorData(shape, bytes) + + val file = Files.createTempFile("bpt-q6k", ".bin") + try { + Files.write(file, bytes) + val storage = MappedBufferStorage.map(file, 0, bytes.size.toLong()) + val buf = BufferPackedTensorData(shape, storage, TensorEncoding.Q6_K) + assertContentEquals(heap.toFloatArray(), buf.toFloatArray()) + assertEquals(heap.get(1, 300).toFloat(), buf.get(1, 300)) + } finally { + Files.deleteIfExists(file) + } + } + + @Test + fun packed_view_is_row_major_over_the_same_storage() { + val shape = Shape(2, 256) + val bytes = randomPayload(2, 144, intArrayOf(0, 2), seed = 3) + val storage = directStorage(bytes) + val buf = BufferPackedTensorData(shape, storage, TensorEncoding.Q4_K) + + val view = buf.packedView + assertEquals(BlockOrder.ROW_MAJOR, view.layout.blockOrder) + assertTrue(view.storage === storage, "packedView must borrow the off-heap storage, not copy") + // rule 4: view.get decodes — spot-check against the heap decode + val heap = Q4_KBlockTensorData(shape, bytes) + assertEquals(heap.toFloatArray()[1 * 256 + 17], view.get(1, 17)) + } + + @Test + fun packedData_refuses_a_heap_copy() { + val bytes = randomPayload(1, 144, intArrayOf(0, 2), seed = 1) + val buf = BufferPackedTensorData(Shape(1, 256), directStorage(bytes), TensorEncoding.Q4_K) + assertFailsWith { buf.packedData } + assertFailsWith { buf.set(0, 0, value = 1f) } + } + + @Test + fun size_and_encoding_are_validated() { + val bytes = randomPayload(2, 144, intArrayOf(0, 2), seed = 2) + assertFailsWith { + // storage holds 2 blocks, shape wants 4 + BufferPackedTensorData(Shape(4, 256), directStorage(bytes), TensorEncoding.Q4_K) + } + assertFailsWith { + BufferPackedTensorData(Shape(2, 256), directStorage(bytes), TensorEncoding.Q8_0) + } + } +}