[michalharakal]

Against engine 0.63.0, transformers-only and additive (no API removed; apiCheck passes for lfm2 and klfm2).

Changes

Lfm2DecodeStep (llm-inference:lfm2, common). The fixed-capacity single-token step of Lfm2CachedTensorNetwork, from an already embedded token and with an output head of its own. One graph serves every position up to the capacity:

• inputs: embedded, cos/sin, the additive mask [1, C+1] and the one-hot write_mask [1, C, 1];
• state in and out: history_<l> and key_<l>/value_<l>.

The head can be a Q4_0 copy of the tied embedding, and then it folds. The tied tensor itself also serves the lookup, so with two consumers the fusion pass skips it. The layer loop and cache bookkeeping move into Lfm2CachedLayers, which Lfm2CachedTensorNetwork now shares. That network's behaviour is unchanged and its tests pass.

Lowerings for the step:

• ShortConvLowering.MULTIPLY_SUM computes the single-token short convolution as sum(window * taps) instead of a grouped conv1d. The result is numerically equal.
• The step's attention is one batched product over all heads by default (AttentionLayoutPolicy.CURRENT_BATCHED). That's 12 attention matmuls instead of 192 for LFM2.5-350M, and 22 % faster on the NPU.

DecodeStepExport.exportPorts(inputs, dtype) { ports, ctx -> results } (llm-core, jvm): the decode-step export over arbitrary named ports, for decoders whose state is more than K/V. export is now built on it, with the same output.

Export:

• exportLfm2DecodeStep traces the step with packed weights kept and folded.
• loadLfm2Model(..., weightFormFor) with lfm2ProjectionsAs(Q8_0) loads the projections requantized. The old loadLfm2Model signature stays.
• Lfm2DecodeStepExportCli (Gradle exportLfm2DecodeStep) writes decode_step.mlir, io.txt (the port order), embed.bf16.bin (the host's embedding rows) and export.json from a verified local checkpoint. It fails if any packed matmul is left unfolded.

External bf16 parameters, for a GPU (--parameters external-bf16, commit 844e0de):

• traceLfm2DecodeStepExternal traces the dense step with every weight as a model-scope parameter.
• writeLfm2Bf16Parameters writes the weights as BF16 safetensors and types the parameter globals bf16, with a convert to f32 where they load.
• The mode also writes io.txt, export.json and the embedding rows.
• Compiled with IREE 3.11 for an arm32 Mali GPU (Vulkan, valhall4): 152 ms per step, logits cosine 0.99991 to FP32, same top-1. There the weights are bf16 and the compute is f32.

Results, LFM2.5-350M (Q8_0 projections, Q4_0 head, capacity 64)

Folded block-quant matmuls	93 of 93
Host, quantized step vs FP32 model	logits cosine 0.997–0.999, the same top-1 at all 10 positions; greedy, the first 14 of 16 tokens equal
SL2610, Torq 2.3.0 (bf16 retype)	one NPU dispatch; 533 ms per step on the NPU, 541 ms in-process with upload and download; module load ≈ 0.6 s
SL2610 vs FP32	golden step logits cosine 0.996, same top-1, 9 of the top 10 shared; greedy, 2 of 2 and 12 of 16 tokens equal on two chat prompts

With every projection as Q4_0, the Torq 2.3.0 compiler fails with "failed to allocate LRAM addresses". Q8_0 projections compile.

Tests

• lfm2 jvmTest: 47 tests, 0 failed. Also new: the external export serves every weight from a bf16 archive. New in Lfm2DecodeStepTest:
  • the step equals the full model at every position, for both convolution lowerings and both attention layouts;
  • multiply-sum equals the grouped conv;
  • the export folds every projection and the head (25 on a small config) and has no stablehlo.convolution;
  • three gated real-checkpoint tests (LFM2_MODEL_DIR): Q8_0/Q4_0 against FP32, a golden step for a device (LFM2_GOLDEN_DIR), and a greedy reference (LFM2_PROMPT_IDS).
• llm-core: exportPortsNamesArbitraryStateAndRejectsAnAliasedResult; the existing DecoderKvModelTest export tests are unchanged.
• apiCheck: lfm2 and klfm2. klfm2 jvmTest passes.

🤖 Generated with Claude Code