Repository navigation
FunctionGemma export: convert MLIR-text rewrites to StableHLO/DAG transforms #248
Description
Activity
- added a commit that references this issue
on Aug 11, 2026 Noting for scope-tracking: PR #302 moves the FunctionGemma export (including the bf16/int8 MLIR-text rewrites this issue is about —
rewriteGlobalsToBf16,rewriteGlobalsToInt8, the sentinel-primeGEMMA_SENTINEL_PASTrollback) from:llm-runtime:kgemmainto the new:llm-inference:functiongemmamodule, verbatim (golden-equivalence gate: byte-identical MLIR/safetensors vs. the pre-move export). The rewrites themselves are untouched — converting them to real StableHLO/DAG transforms stays this issue's scope, just relocated toFunctionGemmaExportHarnessin the new module.Also surfaced while docker-vmfb-parity-testing the moved export (unrelated to the text-rewrite conversion, but possibly of interest here): host-CPU (x64)
iree-run-moduleon the compiledgemma/gemma_with_pastgraphs diverges from the board-verified oracle at the first generated special-vocabulary token, reproducing identically in bf16 and FP32 and independent of the KV-cache loop — looks like a stock-IREE llvm-cpu / embedding-gather behavior for high (added-vocabulary) token indices that the board's aarch64 NEON Torq-fork path doesn't hit. Pre-existing, not a regression from the move. Diagnostic trail is inFunctionGemmaVmfbParityTest's class doc in the new module.
The FunctionGemma KV/int8 decode graphs are DSL/DAG-authored, but the graph post-processing in
FunctionGemmaExport.ktis done by regex rewrites of the emitted StableHLO text, not graph passes.Convert each to the SKaiNET way (a StableHLO/DAG transform):
rewriteGlobalsToBf16) → a DAG/StableHLO dtype transform (leverageDtypeForwardPropagationPass).rewriteGlobalsToInt8injecting dequant as text) → a real graph quantization pass; keep the host per-row quant writer, emiti8→f32 × scaleas graph nodes.7919→x?xstring relax (most fragile: a magic prime that must not collide with a real dim/SSA id) → proper dynamic-shape tracing.GemmaModel.forwardWithPastattnWithPast(hand-wired single-token attention) → an MHA-with-past module forward.refsForpositional sub-module resolution → typedHybridTransformerBlockfield access.Context: PR #245. Tracker: SKaiNET-embedded
sl2610-function-calling/docs/GEMMA-KV-INT8.md.