Skip to content

FunctionGemma export: convert MLIR-text rewrites to StableHLO/DAG transforms #248

Description

@michalharakal

The FunctionGemma KV/int8 decode graphs are DSL/DAG-authored, but the graph post-processing in
FunctionGemmaExport.kt is done by regex rewrites of the emitted StableHLO text, not graph passes.
Convert each to the SKaiNET way (a StableHLO/DAG transform):

  • bf16 weight emission (rewriteGlobalsToBf16) → a DAG/StableHLO dtype transform (leverage DtypeForwardPropagationPass).
  • int8 quant (rewriteGlobalsToInt8 injecting dequant as text) → a real graph quantization pass; keep the host per-row quant writer, emit i8→f32 × scale as graph nodes.
  • Dynamic KV dims — retire the sentinel-prime 7919x?x string relax (most fragile: a magic prime that must not collide with a real dim/SSA id) → proper dynamic-shape tracing.
  • GemmaModel.forwardWithPast attnWithPast (hand-wired single-token attention) → an MHA-with-past module forward.
  • refsFor positional sub-module resolution → typed HybridTransformerBlock field access.

Context: PR #245. Tracker: SKaiNET-embedded sl2610-function-calling/docs/GEMMA-KV-INT8.md.

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