Repository navigation
feat(gemma): export one FunctionGemma decode step for the Torq NPU - #529
Merged
Merged
Conversation
GemmaModel.forwardEmbeddedWithPast runs the decode step from an already embedded token, so the embedding lookup can stay on the host, and takes an optional lmHeadWeight to project the logits with a different (e.g. re-quantized) copy of the tied embedding. forwardWithPast delegates to it. TorqDecodeStepExportTest (gated on TORQ_DECODE_OUT and GEMMA_GGUF) traces one step with a static past of 16 positions, keeps the Q8_0 projections packed, folds all 127 of them with BlockQuantMatmulFusionPass (the vocabulary projection as a Q4_0 copy of the embedding) and writes the StableHLO module plus the argument/result roles. Converted to bf16, it compiles with Torq 2.3.0 to one dispatch on the mixed_quant_matmul kernel and runs on an SL2610 at 162 ms/step with the same top-1 as the f32 reference. Needs a SKaiNET release with SKaiNET-developers/SKaiNET#1478 and #1479. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
GemmaModel and DecoderKvModel each had their own RopeCosSin, ChunkContext and prefill / with-past output classes. They now share llm-core's RopeTables, RopeCosSin (tables per RoPE base: global, plus sliding for Gemma 3), ChunkContext (rope plus a mask per layer type), KvStepOutput and buildRopeCosSin, all generic over the model's DType; both models route per layer by MultiHeadAttention.slidingWindow. The old call shapes RopeCosSin(cos, sin) and RopeCosSin(cosG, sinG, cosS, sinS) still work. DecoderKvModel gains forwardEmbeddedWithPast like GemmaModel, and the decode-step export moves out of the FunctionGemma test into llm-core's DecodeStepExport: it binds the graph inputs to the trace, keeps packed weights, runs BlockQuantMatmulFusionPass and names every argument and result. TorqDecodeStepExportTest keeps only the Gemma-specific parts; its module is unchanged apart from SSA names. :llm-inference:gemma now exposes :llm-core as api, since its public with-past methods use these types. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
0.63.0 ships the block-quant StableHLO export (SKaiNET#1478, #1479) that DecodeStepExport and TorqDecodeStepExportTest build on. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On engine 0.63.0 the scoped and unscoped decode of the tiny llama model differ in 12 of 192 logits over the first two steps, by at most 7.5e-9: one ulp of the step's largest logit, but 16 ulps of a logit near zero (0.0066), which tripped the per-logit four-ulp pin. The accumulation-order rounding is on the scale of the step's activations, so the tolerance is now four ulp of the step's largest |logit|. Engine 0.62.0 passes either way. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Against SKaiNET 0.63.0, which ships SKaiNET-developers/SKaiNET#1478 and SKaiNET-developers/SKaiNET#1479 (
PackedConstantHandling.KEEP_PACKED,BlockQuantMatmulFusionPass, off-heap block reads).Changes
Shared with-past types (
llm-core,sk.ainet.lang.nn.dsl.decoder, generic overT : DType):RopeTables,RopeCosSin(tables per RoPE base:global, andslidingfor Gemma 3),ChunkContext(ropeplus a mask per layer type),KvStepOutput, andbuildRopeCosSin(attention, position, dtype, ctx, rows).GemmaModelandDecoderKvModel. Both models route per layer byMultiHeadAttention.slidingWindow.RopeCosSin(cos, sin)andRopeCosSin(cosG, sinG, cosS, sinS)keep the old call shapes.ChunkContextnow takes aRopeCosSin.:llm-inference:gemmaexposes:llm-coreasapi.forwardEmbeddedWithPast(embedded, rope, selfKIn, selfVIn, ctx, lmHeadWeight = null)onGemmaModelandDecoderKvModel: the decode step from an already embedded token, so the embedding lookup can stay on the host.lmHeadWeightoptionally projects with a replacement output weight.forwardWithPastdelegates to it; behaviour is unchanged.DecodeStepExport(llm-core, jvm): traces one single-token step of any KV-cache decoder with a static past.x * xrun untraced and become a constant.BlockQuantMatmulFusionPass.embedded,cosGlobal/sinGlobal,cosSliding/sinSliding,l<i>.k/l<i>.v,logits.TorqDecodeStepExportTest(functiongemma, gated onTORQ_DECODE_OUTandGEMMA_GGUF): FunctionGemma throughDecodeStepExport, with the vocabulary projection as a Q4_0 copy of the embedding (the Torq compiler hangs on the Q8_0 one). All 127 projections fold. The module is identical to the pre-refactor export apart from SSA names.Build: SKaiNET 0.63.0 (
gradle/libs.versions.toml).ForwardScopeSteadyStateTest(llama): the tolerance is now 4 ulp of each step's largest logit instead of a fixed ulp count of each value. Engine 0.63.0 changes the accumulation order, so near-zero logits move by many of their own ulps (0.0066 moves by 7.5e-9, 16 of its ulps but under 1 ulp of the largest logit). This is needed for any move to 0.63.0, independent of the rest of the PR.Result on an SL2610 (Torq 2.3.0)
After a bf16 conversion of the module (Torq's f32 host fallback returns NaN):
mixed_quant_matmultiles and no generic matmulTests
assemble allTests) green on 0.63.0; localallTests --continuepasses.llm-corejvmTest: 154 tests, 0 failed. New:forwardEmbeddedWithPastMatchesForwardWithPast,decodeStepExportNamesEveryPort, and a check that a single-base model has no separate sliding tables.llm-coreandgemma.FunctionGemmaWithPastCpuTest.with_past_two_graph_loop_matches_oracle(kgemma, real checkpoint) also fails on plain develop against SKaiNET 0.62.0.Not in this PR
A position-independent step (fixed-size cache with a mask), a bf16 export option instead of the text conversion, and app integration.
🤖 Generated with Claude Code