Skip to content

feat/gemma-token-embd-rowdequant needs redesign against transformer-core #234

Description

@michalharakal

feat/gemma-token-embd-rowdequant (021cb11) keeps Gemma's token_embd table packed instead of dequantising it wholesale to FP32, so Embedding dequantises only the row per token at decode time. The claimed win is real and worth having: FunctionGemma Q8_0 token_embd ~178 MB packed vs ~0.67 GB FP32 — per the code comments, "the difference between completing decode and OOMing the 1.9 GB board".

The branch cannot be merged as-is. It was cut 113 commits ago, before the transformer-core extraction, and develop has moved under it.

1. Merging it silently drops the optimization

llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/layers/Embedding.kt no longer exists on develop — commit 4042af0 ("Extract transformer-core: NN primitives reusable on all targets incl. androidNative") relocated it to transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/layers/Embedding.kt.

The branch edits the old llm-core path (+44 lines adding the row-dequant gather). A simulated merge (git merge-tree) produces a tree whose transformer-core/.../Embedding.kt contains zero occurrences of forwardRowDequant, with no llm-core copy. The entire row-dequant path evaporates and RowDequantSource.kt is left behind as dead code. The only conflict git reports is an unrelated one in llm-core.api — the actual loss is silent.

2. The design can't compile against develop

The branch puts RowDequantSource in llm-core and has Embedding reference it. On develop, Embedding now lives in transformer-core, and llm-core/build.gradle.kts declares api(project(":transformer-core")) while transformer-core has no project dependencies. So transformer-core cannot reference an interface in llm-core — that's a dependency cycle.

Fix: relocate RowDequantSource into transformer-core next to Embedding, then re-apply the gather path there.

3. Breaking public API (call it out in the release notes)

sk.ainet.models.gemma.RowDequantSource is removed and re-homed as sk.ainet.lang.nn.RowDequantSource; GemmaPerLayerTokenEmbedTensorData and SafeTensorsPerLayerTokenEmbedTensorData change declared supertypes. Binary-incompatible for external consumers.

Other gaps

  • No tests. The diff adds none, and the commit message contains no verification evidence, benchmarks, or numbers — the quantitative claims live only in source comments.
  • Self-described as WIP: "Add initial support ... Evalute later adding raw packe support in upstream".

Recommend treating this as a rewrite on top of current develop rather than a rebase.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions