Skip to content

functiongemma export: write externals from any BufferHandle (Owned or Floats) — fixes #405 #414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michalharakal merged 2 commits into from
Sep 4, 2026
Merged

functiongemma export: write externals from any BufferHandle (Owned or Floats) — fixes #405 #414

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6686107
functiongemma export: write externals from any BufferHandle (Owned or…
michalharakal Sep 3, 2026
294fb99
Merge remote-tracking branch 'origin/develop' into functiongemma-expo…
michalharakal Sep 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 72 additions & 24 deletions ...ctiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ import sk.ainet.lang.tensor.data.Bf16TensorData
import sk.ainet.lang.tensor.data.TensorData
import sk.ainet.lang.tensor.ops.VoidTensorOps
import sk.ainet.lang.tensor.storage.BufferHandle
import sk.ainet.lang.tensor.storage.DefaultBufferResolver
import sk.ainet.lang.types.FP32
import sk.ainet.models.gemma.GEMMA_DEQUANTIZE_ALL
import sk.ainet.models.gemma.GemmaWeightLoader
Expand Down Expand Up @@ -66,8 +65,6 @@ import java.nio.ByteOrder
public object FunctionGemmaExportHarness {

/** Little-endian bytes of an external parameter, whatever `BufferHandle` the engine handed over (#420). */
private fun bytesOf(h: BufferHandle): ByteArray = DefaultBufferResolver().resolve(h).use { it.readAllBytes() }

public data class RedecodeResult(
val mlirPath: String,
val safetensorsPath: String,
Expand Down Expand Up @@ -463,27 +460,79 @@ public object FunctionGemmaExportHarness {
os.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(headerBytes.size.toLong()).array())
os.write(headerBytes)
for (e in ext) {
val data = bytesOf(e.source)
if (bf16) {
val base = 0
val n = data.size / 4
val obuf = ByteArray(n * 2)
for (j in 0 until n) {
val o = base + j * 4
val fb = (data[o].toInt() and 0xFF) or
((data[o + 1].toInt() and 0xFF) shl 8) or
((data[o + 2].toInt() and 0xFF) shl 16) or
((data[o + 3].toInt() and 0xFF) shl 24)
val bf = Bf16TensorData.floatToBf16Bits(Float.fromBits(fb)) // truncation = core parity
obuf[j * 2] = (bf and 0xFF).toByte()
obuf[j * 2 + 1] = ((bf ushr 8) and 0xFF).toByte()
}
os.write(obuf)
} else {
os.write(data)
writeHandle(os, e.source, bf16)
}
}
}

/**
* Streams one external's f32 payload as raw f32 or truncating bf16 (= core parity), in 1 MiB
* chunks. Dispatches on the [BufferHandle] subtype: the 0.53.0 engine loader delivers large
* constants (the 262144x640 tied embedding) as [BufferHandle.Floats], not [BufferHandle.Owned] —
* the cast that #396 removed from the Gemma 3n harness was still here (#405).
*/
private fun writeHandle(os: java.io.OutputStream, src: BufferHandle, bf16: Boolean) {
val n: Int
val floatAt: (Int) -> Float
when (src) {
is BufferHandle.Owned -> {
val data = src.data
val base = src.offset
n = (src.sizeInBytes / 4).toInt()
if (!bf16) {
os.write(data, base, n * 4)
return
}
floatAt = { j ->
val o = base + j * 4
Float.fromBits(
(data[o].toInt() and 0xFF) or ((data[o + 1].toInt() and 0xFF) shl 8) or
((data[o + 2].toInt() and 0xFF) shl 16) or ((data[o + 3].toInt() and 0xFF) shl 24),
)
}
}
is BufferHandle.Floats -> {
val f = src.data
n = f.size
floatAt = { j -> f[j] }
}
else -> error("unsupported BufferHandle ${src::class.simpleName} (${src.sizeInBytes} B)")
}
val bpe = if (bf16) 2 else 4
val chunk = (1 shl 20) / bpe
val buf = ByteArray(chunk * bpe)
var j = 0
while (j < n) {
val m = minOf(chunk, n - j)
if (bf16) {
for (k in 0 until m) {
val bf = Bf16TensorData.floatToBf16Bits(floatAt(j + k)) // truncation = core parity
buf[k * 2] = (bf and 0xFF).toByte()
buf[k * 2 + 1] = ((bf ushr 8) and 0xFF).toByte()
}
} else {
val bb = ByteBuffer.wrap(buf).order(ByteOrder.LITTLE_ENDIAN)
for (k in 0 until m) bb.putFloat(floatAt(j + k))
}
os.write(buf, 0, m * bpe)
j += m
}
}

/**
* Little-endian f32 bytes + base offset for any [BufferHandle]; a [BufferHandle.Floats] is
* materialised once (the quantizer reads rows by byte offset). See [writeHandle] / #405.
*/
private fun ownedBytes(src: BufferHandle): Pair<ByteArray, Int> = when (src) {
is BufferHandle.Owned -> src.data to src.offset
is BufferHandle.Floats -> {
val f = src.data
val b = ByteArray(f.size * 4)
val bb = ByteBuffer.wrap(b).order(ByteOrder.LITTLE_ENDIAN)
for (x in f) bb.putFloat(x)
b to 0
}
else -> error("unsupported BufferHandle ${src::class.simpleName} (${src.sizeInBytes} B)")
}

private fun voidF32(shape: Shape): sk.ainet.lang.tensor.Tensor<FP32, Float> =
Expand Down Expand Up @@ -617,8 +666,7 @@ public object FunctionGemmaExportHarness {
os.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(headerBytes.size.toLong()).array())
os.write(headerBytes)
for (e in ext) {
val data = bytesOf(e.source)
val base = 0
val (data, base) = ownedBytes(e.source)
val qs = quant[e.key]
if (qs != null) {
val (rows, cols) = qs
Expand All @@ -640,7 +688,7 @@ public object FunctionGemmaExportHarness {
os.write(q)
os.write(sb)
} else {
val n = data.size / 4
val n = e.source.sizeInBytes.toInt() / 4
val ob = ByteArray(n * 2)
for (j in 0 until n) {
val bf = Bf16TensorData.floatToBf16Bits(leF32(data, base + j * 4))
Expand Down
Loading