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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions