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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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