Gemma 4 engine gaps: gemma4 registries, KernelDispatch self-heal, and dense FP32 matmul (fixes #1220) - #1221
Merged
Conversation
…n, observable dispatch fallback Four fixes surfaced by the Gemma 4 E2B GGUF debugging session (SKaiNET-transformers#325/#330 arc): - TokenizerFactory: accept 'gemma4' as tokenizer.ggml.model (SentencePiece + specials wrap) — the allowlist threw UnsupportedTokenizerException on real Gemma 4 GGUFs even though the pipeline handles them. - TokenizerFactory: read tokenizer.ggml.token_type via toIntFlexible — GGUF UINT32 arrays surface as kotlin.UInt (not Number), so the old 'as? Number' filter silently dropped every special token on such files. - SpecialTokenSplitter: override decodeToken to delegate to the base's decodeToken — the interface default routed through the batch decode path, re-enabling leading-space stripping and losing word-boundary spaces in per-token streaming for every SentencePiece GGUF with specials. - KernelDispatch: settable defaultSink + a one-time loud warning when the decoding reference matmul serves a blocked weight — the ~1000x fallback used to be invisible because production call sites pass no TraceSink. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
…Pack SPI KernelDispatch was the only registry that did not populate itself. KernelRegistry self-heals (DefaultCpuOpsJvm.ensureKernelProviders installs providers lazily on first use), but every dispatch consumer had to remember KernelPacks.install() + FfmRowMajorKernelPack.install() before its first forward pass. Forgetting it is invisible: the decoding reference kernel is correct for every format, just about a thousand times slower. It was forgotten repeatedly in practice — by an application entry point (SKaiNET-transformers' kgemma CLI) and by a diagnostic harness, where it turned a ten-forward-pass run into 23 minutes of CPU. - New ViewKernelPack SPI in skainet-backend-api: an installable set of ViewKernels, the view-keyed sibling of KernelProvider. Discovery is ServiceLoader-based on JVM and Android and manual elsewhere, the same split KernelProvider already documents (expect/actual across all six source sets). - KernelDispatch.ensureInstalled(), called at the top of matmul(): when nothing is registered, discover providers, then KernelPacks.install(), then every discovered pack. Providers must come first — KernelPacks.install() derives its kernels from KernelRegistry.bestAvailable(), which is null on an empty registry, so the reverse order silently installs only the reference kernel (measured: 8 of 17 kernels). - FfmRowMajorKernelPackFactory (native-cpu) and JniMappedKernelPackFactory (jni-cpu) plus their META-INF/services entries, so the row-major kernels that serve mapped GGUF weights zero-copy are found without a consumer call. - Explicit registration still wins: ensureInstalled() is a no-op once the table is non-empty, so a curated kernel set is never silently widened. Verified end-to-end: with the downstream bootstrap removed entirely, a Gemma 4 E2B Q4_K_M parity run loads and does 10 forward passes in 8.7s and produces the same output as with an explicit bootstrap. Docs: new explanation/kernel-selection.adoc covering both registries, the SPI and its registration/discovery rules, the selection algorithm (including the adapter and reference fallbacks), why a kernel declines, and how to diagnose which kernel ran. Cross-linked from explanation/eager-execution.adoc, which covers the provider tier only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
…st kernels can read it matmulWeightTransposed materialized Wt on EVERY call for a dense weight, while the packed path has cached its relayout 'once per weight instead of once per call' since #1096. transpose() cannot return a view here — the FP32 kernels decline any weight whose strides[0] != 1, and no view of a row-major [out, in] buffer has that layout — so each call paid a full cache-hostile scatter over every element. Three changes, all in that one asymmetry: - cache Wt per weight, keyed on TensorData identity (not a FloatArray: a weight staged by the MemorySegment factory is not FloatArray-backed, and that is the case that hurts most — transpose then falls to its generic per-element fallback); - materialize the cached copy on the HEAP rather than through dataFactory, because the vectorized FP32 kernels take heap FloatArray operands only, so a segment-backed transpose is condemned to the reference kernel afterwards; - widen the JVM activation->heap bridge from packed weights to all weights, for the same reason on the other operand. Measured on Gemma 4 E2B Q4_K_M (its per-layer-embedding projections are the only dense-FP32 weights in an otherwise Q4_K model, transposed 70x per token): decode 443 -> 272 ms/token (2.26 -> 3.68 tok/s) prefill 432 -> 231 ms/token (2.31 -> 4.32 tok/s) Correctness unchanged: the Gemma 4 golden-token gate, its GGUF smoke test, FunctionGemma's tokenizer/template/e2e gates and the engine backend suites all pass. Note the answers move in the last ulp — a vectorized kernel accumulates in a different order than the scalar loop it replaces. Still open: even heap-on-heap, a dense FP32 projection runs ~0.3 GFLOP/s against ~29 for the packed Q4_K path beside it, so those two matmuls remain 73% of decode. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
… is overhead-bound there The FP32 SPI kernel is tile-blocked, and at decode sizes its per-call setup is essentially the whole cost. Measured (Fp32GemvShapeBench, k=1536 n=256): m= 1 1.002 ms/call 0.78 GFLOP/s m= 8 1.075 ms/call 5.85 GFLOP/s m= 2 1.028 ms/call 1.53 GFLOP/s m=16 1.260 ms/call 9.98 GFLOP/s m= 4 0.992 ms/call 3.17 GFLOP/s m=32 1.600 ms/call 15.73 GFLOP/s Eight times the arithmetic for 8% more time: below m~32 this is a fixed ~1ms overhead, not a throughput curve. Decode is m=1 by construction, so a model whose checkpoint ships dense FP32 tensors pays it per layer per token — Gemma 4 E2B has 70 such projections (its per-layer-embedding gate and proj are F32 in the GGUF while everything around them is Q4_K) and spent 73% of decode there. Small shapes now go through a direct i-p-j loop instead: contiguous in b and out, one pass, no setup. The threshold (4M MACs) covers decode-step projections and leaves prefill batches and the vocab matmul to the tiled kernel, which is already ~16 GFLOP/s at m=32 and pulling away. m=1 after: 0.125 ms/call, 6.30 GFLOP/s — 8x End to end on Gemma 4 E2B Q4_K_M, cumulative with the transpose-cache commit: decode 443 -> 189 ms/token (2.26 -> 5.29 tok/s) 2.3x prefill 432 -> 186 ms/token (2.31 -> 5.37 tok/s) 2.3x the two PLE projections: 1518 -> 49 ms over 8 tokens 31x and the decode profile is balanced again: lm_head 29%, attention 28%, RoPE 14%, PLE 15%, where PLE alone used to be 73%. Gates green: Gemma 4 golden-token and GGUF smoke, FunctionGemma tokenizer/template /e2e, tokenizer parity, and the engine backend suites. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
…ernel Three paths measured on an Apple M4 at k=1536 n=256 (Fp32GemvShapeBench), so the next person does not re-derive them: m=1 heap + direct loop 0.122 ms 6.44 GFLOP/s <- what we do now m=1 heap + native FFM kernel 1.002 ms 0.78 GFLOP/s <- copy-bound m=1 off-heap, zero copy 0.943 ms 0.83 GFLOP/s <- kernel-bound The native kernel copies both operands off-heap per call because a JDK 21 downcall cannot address heap memory, and that copy is proportional to the weight rather than to the work: 1.5 MB whatever m is. Keeping the weight off-heap instead — it is read-only for the life of the process, so this is the obvious escape — removes the copy and is still slower, because the segment kernel that then serves it is weaker than the heap one by more than the copy costs. So residency is not the lever. There are two independent gaps: this kernel cannot see heap memory (fixed on JDK 22+ by Linker.Option.critical(true), noted at the call site), and matmulFloatBlockedMemSeg is slow. Either is worth real throughput — the packed Q4_K path reaches ~29 GFLOP/s on the same machine. Also recorded: two things that measured as no-ops and should not be retried — reusing the Arena and its segments across calls (1.025 vs 1.002 ms at m=1; the allocation was never the cost) and reordering the C loops (already i-p-j). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
…hape Three benchmarks that turn 'the dense path is slow' into per-kernel numbers. m=1, k=1536, n=256 (a Gemma 4 E2B per-layer-embedding projection), Apple M4: native-ffm bf16 0.055 ms 14.33 GFLOP/s <- fastest kotlin-direct fp32 0.073 ms 10.78 GFLOP/s <- what DefaultCpuOpsJvm now uses native-ffm fp32 0.105 ms 7.48 GFLOP/s native-ffm fp16 0.121 ms 6.51 GFLOP/s scalar fp32 0.253 ms 3.10 GFLOP/s panama bf16 0.466 ms 1.69 GFLOP/s scalar fp16 0.886 ms 0.89 GFLOP/s panama fp32 0.969 ms 0.81 GFLOP/s Corrections to earlier guesses in this area, all measured: - The native FFM kernel is NOT the slow one. Its heap->off-heap copy, the obvious suspect, is 0.029 ms for the same 1.5 MB (54 GB/s, HeapToSegmentCopyBench) — 3% of a call, not the cost. - PanamaVectorMatmulKernel is slower than the SCALAR kernel at this shape, by 3.8x (0.969 vs 0.253 ms). It is the fallback wherever the native library is absent, so that is a live gap, not a curiosity. - Narrower weight storage does help, but BF16 not FP16: bf16 is 1.9x faster than fp32 (halved weight traffic, decode is a bit-shift), while fp16 is *slower* than fp32 because binary16 needs exponent rebiasing — exactly what the SPI doc says. Accumulation stays FP32 in both, so this is memory traffic, not math precision. - Off-heap residency does not help either: the segment kernel that serves it is slower than the heap one by more than the copy costs (0.943 vs 0.122 ms). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
…— add a GEMV path The tiled path has two fixed costs, both O(n*k) and both independent of m: it packs B transposed into a fresh n*k FloatArray on every call (a scatter, 1.5 MB for the shape below), and it does a horizontal reduceLanes per output cell per K-tile. Neither buys anything when there are no rows to amortize them over — and a decode step is m=1 by construction. Measured at k=1536 n=256 on an Apple M4, this kernel was losing to the scalar one it exists to beat: 0.969 ms vs 0.253. At or below 8 rows it now accumulates by outer product instead — out[i,:] += a[i,p] * b[p,:] — which touches B once, contiguously, keeps each lane on its own output column for the whole k loop, and never reduces across lanes: m before after native-ffm scalar 1 0.969 ms 0.81 0.058 ms 13.63 7.65 3.05 GFLOP/s 4 — 0.230 ms 13.70 18.50 3.07 8 — 0.463 ms 13.58 21.41 3.02 16 1.332 ms 9.45 (tiled, unchanged) 25.53 3.06 32 1.647 ms 15.28 (tiled, unchanged) 27.69 3.10 16.7x at m=1, and it becomes the fastest FP32 kernel there — ahead of native-ffm, whose heap->off-heap copy costs more than the arithmetic at that shape. The tiled path is untouched above 8 rows, where it is what the shape wants. This matters most where there is no native library: Panama is then the best available provider, so every platform without the FFM backend was running its dense FP32 matmuls ~17x slower than necessary. Accumulation for a given output stays in ascending p, matching the scalar kernel's order rather than the tiled path's split-by-K-tile order; the native/Panama parity suite and the downstream Gemma 4 and FunctionGemma gates all pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
…spose cache private Two CI breaks from the ViewKernelPack commit: - :skainet-backend-api:compileAndroidMain failed because the Android actual delegated to KernelServiceLoader, which lives in jvmMain and is not visible from androidMain. Inlines the same two steps it performs — ServiceLoader.load, then KernelRegistry.register, which sorts by priority on insertion. - jvmApiCheck failed because transposedDenseWeight was declared protected, putting an implementation detail into skainet-backend-cpu's public API. It is only called from within DefaultCpuOps, so it is private now and the .api dump is unchanged. Verified locally: compileAndroidMain, apiCheck across the repo, and the backend-api / backend-cpu / backend-native-cpu jvm suites. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
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
1 participant
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.
Fixes #1220 — the engine half of getting Gemma 4 E2B Q4_K_M to generate correct text at a usable speed. Seven fixes in three groups; the issue has the full investigation, this is what changed and how it was checked.
Registries and tokenizer (
skainet-io)TokenizerFactory: acceptgemma4— an ordinary SentencePiece vocab with CONTROL/USER_DEFINED specials, which the factory refused by model-string allowlist.TokenizerFactory: readtokenizer.ggml.token_typeviatoIntFlexible. GGUF UINT32 arrays arrive askotlin.UInt, not aNumber, so the oldas? Numberfilter silently dropped every special token on such files. The helper existed for this and wasn't wired here.SpecialTokenSplitter.decodeToken: delegate to the base tokenizer instead of inheriting the batch-decode default, which re-enabled leading-space stripping and cost word-boundary spaces in per-token streaming for every SentencePiece GGUF with specials.ModelArchitecture.ggufIdMap:"gemma4"→GEMMA.KernelDispatchself-heals (skainet-backend-api,-native-cpu,-jni-cpu)New
ViewKernelPackSPI — the view-keyed sibling ofKernelProvider— discovered byServiceLoaderon JVM/Android and installed manually elsewhere, the same splitKernelProviderdocuments.KernelDispatch.ensureInstalled()runs at the top ofmatmul: discover providers, thenKernelPacks.install(), then discovered packs. Order matters —KernelPacks.install()derives frombestAvailable(), which is null on an empty registry, so the reverse order installs 8 of 17 kernels.Explicit registration still wins:
ensureInstalled()is a no-op once the table is non-empty, so a curated set is never silently widened.Verified by deleting the downstream bootstrap entirely: an E2B parity run then loads and does 10 forward passes in 8.7s, having previously needed 23 minutes of CPU when the bootstrap was forgotten.
Dense FP32 matmul (
skainet-backend-cpu)Wᵀper weight instead of rebuilding it per call, and materialize it on the heap where the vectorized kernels can read it.m=1by construction.PanamaVectorMatmulKernelgains a GEMV path. It was slower than the scalar kernel atm=1(0.969 vs 0.253 ms) because it packs Bᵀ into a freshn*karray per call and does a horizontalreduceLanesper output cell per K-tile — bothO(n*k), both pointless with one row. At or below 8 rows it now accumulates by outer product: touches B once, contiguously, each lane owning one output column, no lane reductions. The tiled path is untouched above that.16.7x, and it becomes the fastest FP32 kernel at that shape. This matters most where there is no native library — Panama is then the best available provider.
End to end on Gemma 4 E2B: decode 443 → 189 ms/token, prefill 432 → 186 (2.3x each).
Checks
skainet-backend-api,-cpu,-native-cpu,skainet-io-core.p, matching the scalar reference rather than the tiled path's split-by-K-tile order.TransposedWeightViewTestcompared the two spellings bit-for-bit, which held only while both took the same scalar path. Now compared within tolerance, with the reason recorded.Reproducibility
Fp32GemvShapeBench,Fp32KernelRaceBenchandHeapToSegmentCopyBenchare committed so the numbers above can be re-run rather than trusted. New Antora pageexplanation/kernel-selection.adocdocuments both registries, the SPI and its discovery rules, the selection algorithm with its adapter and reference fallbacks, why a kernel declines, and how to tell which one ran.Kernel doc comments also record the paths that were measured and rejected — off-heap residency, arena reuse, C-loop reordering — so they aren't retried.