FunctionGemma redecode contract: LM head + argmax run over all SEQ positions (36 % of every step, 1 GB scratch at SEQ 1024 in a 32-bit process); parameter archives are per-SEQ

[michalharakal]

Summary

The exported redecode graph (func @gemma(tensor<1xSEQxi32>) -> tensor<SEQxi32>) computes the LM head for every position and then materialises three more SEQ×262144 tensors for the argmax, although the caller only ever reads one position per step:

%v3325 = stablehlo.dot_general %hidden, %embT : (tensor<1x64x640xf32>, tensor<640x262144xf32>) -> tensor<1x64x262144xf32>
%v3327 = stablehlo.reduce(... maximum ...) -> tensor<1x64xf32>
%v3329 = stablehlo.compare EQ ... -> tensor<1x64x262144xi1>
%v3330 = stablehlo.iota dim = 2 : tensor<1x64x262144xi32>
%v3332 = stablehlo.select ... -> tensor<1x64x262144xi32>
%v3334 = stablehlo.reduce(... minimum ...) -> tensor<1x64xi32>

(ops.argMax(logits, dim = -1) at FunctionGemmaExportHarness.kt:207.)

Measured consequences (arm32 Android device, IREE 3.11.0, FP32 archive, 13-token prompt)

graph (SEQ 64)	CPU arm32 local-task, 4 threads	Vulkan (Mali, host-side gather)
as exported	10.40 s / step	4.00 s
head sliced to the target position before the dot (hand-patched MLIR)	6.61 s (−36 %)	3.16 s
body only, no head	4.49 s	—

• SEQ 1024 (needed for an 853-token tool catalog prompt) cannot run in a 32-bit process: the step dies on calloc(1, 1076363436) — the 1024×262144 i32 iota/select scratch. With the head sliced, SEQ 1024 runs on both CPU (77.3 s) and GPU (29.9 s).
• Cost is linear in SEQ (SEQ 256 = 4.03× SEQ 64), so every prompt token costs ≈ 0.16 s per generated token on the CPU.
• Parameter archives are per-SEQ: the export bakes SEQ-sized position tables into the parameters (parameter t46 is 32 KB at SEQ 64, 128 KB at SEQ 256; a SEQ 256 module opened against the SEQ 64 .irpa fails with parameter 't46' range out of bounds). A cartridge shipping two SEQ tiers ships two 1 GB archives.

Proposal

1. Add the target position as an input (or define "last non-pad position" semantics) and slice the hidden state before the LM head; return tensor<1xi32> (or keep SEQxi32 with the one computed position, for compatibility).
2. Emit the RoPE/position tables as computed values or as SEQ-independent parameters so one .irpa serves every SEQ.
3. Both are prerequisites for the KV-cache contract (separate issue) and for the Vulkan tier, which is otherwise the fastest path on this device.

Context. Measured on 2026-09-03 while bringing FunctionGemma-270M up as an NLU cartridge on an arm32 Android device (Android 14, armeabi-v7a-only 32-bit process, 4× ARMv8 @ 2.0 GHz, Mali GPU) with SKaiNET 0.53.0 + SKaiNET-transformers 0.53.0 from Maven Central. Harness and raw result files: a local harness (can share on request).