feat(dtype): narrow-float KEEP_NATIVE weights end to end - #263
Merged
Conversation
michalharakal
changed the title
michalharakal
marked this pull request as draft
July 30, 2026 08:55
…F paths Engine 0.38.0 adds the shared narrow-float codec, Fp16DenseTensorData, FP16 matmul kernels and codec-driven matmul dispatch. Take advantage of all of it. - DecoderSafeTensorsLoader gains the F16 KEEP_NATIVE arm that BF16 has had since 0.25.0 — it was missing only for want of a storage type to back it. Covers LLaMA, Qwen and Voxtral, which share this loader. - DecoderGgufWeightLoader accepts a DTypePolicy and keeps F16/BF16 sources in their on-disk layout. The GGUF branches of the network loaders used to construct it without the policy, so withDtypePolicy was ignored on that path entirely; it is now plumbed through. - DTypePolicyValidation takes keepNative: Set<DType> in place of a BF16-only boolean, which could express neither "keeps FP16 but not BF16" nor the empty case. The old overload stays, deprecated. The two narrow formats are resolved independently: Require(BF16) still widens F16 sources and vice versa. Both are 2 bytes per element, so mis-tagging F16 bytes as BF16 would decode to plausible-looking garbage rather than fail. The GGUF packed path swaps the rank-2 shape to [cols, rows] and moves no bytes, mirroring createTensor exactly. GGUF header dims are reversed relative to the logical row-major shape, which is why transposeColumnMajorToRowMajor returns its input untouched; an actual element transpose here would have handed the matmul kernel a silently transposed weight matrix. Gemma and Apertus now reject Require(BF16) instead of accepting a policy their own weight chains never honored and silently disregarding it at load. Callers need Prefer(BF16) until those chains grow a KEEP_NATIVE path. DecoderGgufWeightLoader's constructors gain a trailing defaulted parameter, which changes their JVM descriptors: binary-breaking, source-compatible. Tests: 22 new — 6 on the validation capability model, 5 end-to-end on the SafeTensors path against synthesized files, 11 on the GGUF policy decision and tensor construction (no GGUF writer exists to test that path end to end). 285 green across the touched modules; apiDump regenerated.
Loading tests proved the packed bytes survive and decode correctly, but nothing exercised those tensors once they reached the model. Both narrow formats are 2 bytes per element, so decoding one as the other does not throw — it yields finite, plausible, wrong logits. Synthesize a complete 1-layer LLaMA as SafeTensors, load it twice from the same file (widened vs KEEP_NATIVE), run both through OptimizedLLMRuntime over a 4-token sequence, and compare logits. FP16 and BF16 both match bit-for-bit, including through embedding gather and the RMSNorm weight multiply. Guards against a vacuous pass: - every loaded tensor must be NarrowFloatTensorData, else the comparison would silently be FP32 vs FP32 - weights are round-tripped through the codec before writing, so both sides hold identical values - a codec mix-up case writes F16 bit patterns into a BF16-declared file and asserts the logits move by more than 10x the tolerance Also pin the reason the two paths agree exactly. LlamaRuntime.linearProject calls w.t() on [out, in] weights, transpose has no narrow-float path and widens to a dense FP32 buffer, and chooseQuantizedMatmul only engages for [in, out] weights. The FP16/BF16 kernels are therefore unreachable on this chain: packed weights are widened on every forward, so the saving is at-rest memory only, paid for with a per-token decode and a transpose allocation. Removing the .t() should flip that test and make the parity cases exercise the kernel path they were written for.
Scoping the work to reach the FP16/BF16 kernels needed a number, not a guess: the layout change is only worth doing if the kernels beat the FP32 SGEMM they would replace. Measures three paths at real projection sizes — fp32 dense, narrow weight already in [in, out] so chooseQuantizedMatmul dispatches, and the [out, in] + .t() path production takes today. Skipped unless -Dskainet.bench.narrow=true, so it stays out of CI. The build passes that property through because Gradle does not forward -D to the test JVM. Dispatch is guaranteed by construction rather than assumed: checkDispatchable asserts chooseQuantizedMatmul's own preconditions before timing, and each narrow result is compared against the fp32 baseline so a fast-but-wrong kernel cannot pass unnoticed. Baseline on i7-9750H / OpenJDK 21 recorded in the KDoc. Three results: - BF16 beats FP32 by 1.5-2.1x at every size and batch. Batch-1 matmul is memory-bandwidth bound, so halving the weight bytes roughly halves the time. This is the case for doing the layout work. - FP16 is 2-18x slower, pinned near 0.5 GFLOP/s regardless of shape or batch. Not a layout problem: both Panama kernels fill a scratch array scalar-wise before the vector FMA, but Fp16Codec.decode is a branchy when with a subnormal renormalization loop where the BF16 decode is three integer ops. - The transpose path costs 0.2-4.5 seconds per projection, because the generic transpose walks narrow data element by element through get(). That is per weight, per token. At 8B sizes KEEP_NATIVE is not merely un-accelerated today, it is unusably slow. The earlier forward-parity test could not surface the last point because its model is dim=8, where an elementwise transpose is free.
Weights arrive [out, in] but the narrow matmul dispatch needs [in, out], so Linear.onForward transposes before every matmul. A row-major narrow tensor has no fast transpose — it widened elementwise through boxed get(), measured at 209 ms for a 2048x2048 projection and 4.5 s for 4096x11008, per weight per token. KEEP_NATIVE was slower than not using it at all. Relay matmul weights once at load via the engine's new NarrowFloatInputMajorTensorData (SKaiNET #888), which makes the per-forward transpose a zero-copy view, so the packed weight reaches the narrow kernel. Two kinds of tensor stay row-major. Rank-1 norms are never transposed and the input-major type rejects them outright. The token embedding is gathered by row, not multiplied, so input-major storage would stride exactly the reads it serves. That also covers tied embeddings: when output.weight aliases token_embd it stays row-major and forgoes the transpose win, since one buffer cannot suit both access patterns. Measured on i7-9750H / OpenJDK 21, ffn_up 8B at batch 1, BF16: 4465 ms -> 58 ms, a 77x reduction, and now 1.9x faster than the FP32 SGEMM it replaces. The benchmark gains t()row-maj and t()in-maj columns so the difference stays visible, plus an assertion that the relaid weight is still narrow after t() — without it the new columns could silently re-measure the generic path. The forward-parity logits are no longer bit-identical to the widened path: they diverge by ~7e-9, which is the accumulation-order difference of a genuinely different kernel, and is the evidence the kernel now runs. Tolerance tightened from 1e-4 to 1e-5 accordingly — still three orders of magnitude above the observed delta and three below what a codec mix-up produces. Loader tests updated for the layout split: matmul weights assert input-major and value preservation rather than verbatim bytes, and the verbatim-bytes property is now pinned on the embedding, where it still holds. FP16 remains slower than FP32 here; that is the decode cost tracked in SKaiNET #887 and is unaffected by layout. Prefer BF16 for speed until it lands.
The recorded baseline predates the engine fixes it was written to track, so its headline conclusion is now backwards: it says FP16 is 2-18x slower than FP32, pinned near 0.5 GFLOP/s, blames Fp16Codec.decode's subnormal renormalization loop, and tells the reader to prefer BF16 until #887 lands. FP16 is now 1.5-1.7x faster than FP32 and #887 is merged. Re-measure against engine develop with #888, #887 and the BF16 amortization all merged, and rewrite the conclusions around what the numbers actually show. The decode diagnosis in the old text was the issue's original theory and it was wrong, so the replacement says what the cause turned out to be: NativeKernelProvider carried matmulBf16 but no matmulFp16, so BF16 ran the native FFM kernel at priority 100 while FP16 cascaded to the JVM Panama kernel at 50. Head to head the two Panama kernels are within ~15% of each other. Worth recording, because the failure presented as a slow kernel and was a missing one. Adds a fourth conclusion for the other thing that fell out: both native kernels now read B once per matmul rather than once per row of A, which cut B traffic 16x at m=16 and bought only 9-19%. These kernels are compute-bound on the FMA chain at batch 16, not bandwidth-bound, so further layout work is not where the next win is. Flags the q_proj 1B batch-1 row as noisy rather than quietly re-rolling it: its FP32 sample came out at 5.35 ms against a 3.09-3.22 ms cluster in repeat runs. It is the shortest measurement in the table.
… names compileTestKotlinIosArm64 fails on two backtick test names in DTypePolicyValidationTest: Kotlin/Native rejects "()" and "," in a declaration name because the name has to survive ObjC export. e: Name contains illegal characters: "()" e: Name contains illegal characters: "," The names are commonTest, so every native target compiles them, but the branch never got that far before -- CI died at dependency resolution looking for the unreleased 0.38.0-SNAPSHOT, so `assemble allTests` never reached the iOS compile. Pinning the released engine exposed it. Rename both; the assertions are untouched. Also swept the rest of commonTest for the same pattern -- these two were the only ones. Verified with the exact CI invocation, not just jvmTest: ./gradlew --no-configuration-cache assemble allTests is green.
michalharakal
marked this pull request as ready for review
July 31, 2026 12:46
michalharakal
force-pushed
the
chore/validate-engine-0.38.0-snapshot
branch
from
July 31, 2026 12:46
a877d96 to
161c4b2
Compare
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.
Consumes the engine's narrow-float work end to end and validates it on real model loading paths. This is the downstream half of engine issues #884, #888 and #887.
Rebased onto develop (released
skainet = 0.38.0, which carries the merged narrow-float engine work — PRs #886, #895, #896, #897). Builds and tests against Maven Central directly; the earlier0.38.0-SNAPSHOT+ scoped-mavenLocalvalidation shim is dropped.What it does
KEEP_NATIVE) instead of widening to FP32 at load, halving weight bytes at rest.NarrowFloatInputMajorTensorData, so the per-forward.t()inLinear.onForwardbecomes a zero-copy view. Weights arrive[out, in]but the narrow matmul dispatch needs[in, out]; without the relayout that transpose widened the tensor elementwise through boxedget()— measured at 4.4 s per projection forffn_up8B, per token. That was engine issue #888.output.weightaliasestoken_embdit stays row-major and forgoes the transpose win, since one buffer cannot suit both access patterns.DTypePolicyValidationgains the narrow-float policy surface, with loader wiring across llama/qwen/gemma/apertus/voxtral.Measured
NarrowFloatMatmulBenchmark(jvmTest, opt-in via-Dskainet.bench.narrow=true). i7-9750H (AVX2), OpenJDK 21.0.11, median ms per call, against engine develop with all the above merged:t()in-majmatches the direct column at every size, so the zero-copy view holds. Against the old path that is 4389 ms → 57 ms forffn_upBF16 at batch 1.t()row-majis still measured on purpose: row-major narrow tensors keep the widening behaviour, and only the input-major type may be reinterpreted.q_proj 1Bat batch 1 is the shortest measurement and the noisiest — its FP32 sample came out at 5.35 ms against a 3.09–3.22 ms cluster in repeat runs, so treat that row's ratios with suspicion. The other seven are stable.Tests
KEEP_NATIVEagainst the FP32 path.DTypePolicyValidationTestfor the new policy surface.checkDispatchablemirrorschooseQuantizedMatmul's preconditions before timing, and asserts the relaid weight is still narrow after.t(), so a silent fallback to the generic path cannot be mistaken for a fast kernel.56 tests green in
:llm-inference:llama:jvmTestagainst released 0.38.0 (./gradlew :llm-inference:llama:jvmTest).