feat: shared narrow-float (BF16 + FP16) layer#886
Open
michalharakal wants to merge 4 commits into
Open
Conversation
First stage (B1) of #884. Introduces one 16-bit float layer parameterized by codec, so BF16 and FP16 share an implementation instead of growing two half-built stacks. BF16 had storage, kernels, dispatch and a SafeTensors KEEP_NATIVE path but no trace capture, no elementwise ops and no `convert` target; FP16 had the dtype tag, `convert` and trace capture but no storage, no kernel and no loader. The gaps are complementary, so finishing either alone leaves the other broken. Added: - `NarrowFloatCodec` with `Bf16Codec` and `Fp16Codec` (skainet-lang-core, commonMain — pure integer bit math, no JDK dependency, so every KMP target gets it). Becomes the canonical home for float<->half conversion; the six duplicated private `halfToFloat`/`decodeHalf` copies can now collapse onto it. - `NarrowFloatTensorData` / `NarrowFloatDenseTensorData` — the shared 2-byte dense storage and the recognition surface for narrow-float kernel dispatch. - `Fp16DenseTensorData` — the class the SafeTensors and GGUF loaders already name in their "no FP16 backing yet" errors, which is why `Require(FP16)` had to be rejected outright. Unblocks the loader work (B3). `Bf16Codec.encode` TRUNCATES rather than rounding to nearest. This is contractual, not an oversight: it matches the existing `floatToBf16Bits`, the safetensors writer, and the on-device A/B that qualified bf16 as a bit-exact drop-in for the f16 vmfb. A regression test pins the two implementations together so a future "fix" fails loudly. `Fp16Codec` implements full IEEE 754 binary16 with round-to-nearest-ties-to-even: normals, gradual underflow into subnormals, overflow to +/-Inf past 65504, and NaN payload preservation (never collapsing a NaN into an infinity). Fixed two latent bugs found while wiring this up: - `TraceSession.refOf` had no `BF16::class` arm, so a BF16 tensor traced as FP32. The captured dtype is what the StableHLO converter reads to pick an MLIR element type, so this is precisely why bf16 weights could only be produced by regex-rewriting the emitted MLIR text after the fact (SKaiNET-transformers#248). - `TensorStorageFactory.toTensorData` decoded FLOAT16/BFLOAT16 through the FLOAT32 branch's hard 4-byte reader, yielding half the element count with each "float" assembled from two adjacent narrow elements — silent garbage. Readers already tagged these `Dense(bytesPerElement = 2)`; only the decode side ignored it. Widening to f32 here preserves the method's existing contract; preserving 2-byte storage end-to-end is the loader KEEP_NATIVE path, not this one. `Bf16DenseTensorData` is re-parented onto `NarrowFloatDenseTensorData` and `Bf16TensorData` onto `NarrowFloatTensorData`. Public API is preserved: `floatToBf16Bits`/`bf16BitsToFloat` now delegate to `Bf16Codec`, and the moved members remain resolvable via normal JVM hierarchy lookup. NOTE: `Bf16TensorData` gains `getCodec()` with a default body, so an out-of-tree class implementing that interface and compiled against the old API would hit AbstractMethodError. `Bf16DenseTensorData` is the only implementor in this repo. Tests: 41 new (21 codec, 16 dense storage, 4 storage-decode regression). The pre-existing BF16 suites pass unmodified — Bf16TensorDataTest (12), PanamaVectorBf16MatmulKernelParityTest (10), Bf16MatmulDispatchTest (3); the last of these confirms `is Bf16TensorData` still matches after re-parenting. apiDump regenerated; apiCheck green. Full engine build clean (2133 tasks). Refs #884 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Second stage (B3) of #884. Both loaders can now hand 16-bit weights through in their on-disk 2-bytes-per-element layout instead of widening every one to FP32 at load. Builds on the Fp16DenseTensorData added in B1. SafeTensors: - F16 gains the KEEP_NATIVE arm that BF16 already had. It was absent purely because no Fp16DenseTensorData existed to back it. - `mapPolicyToNarrow(policy, native)` replaces the two hand-written mappings, so one function resolves the policy for either format. - `Bf16LoadPolicy` becomes a typealias to the new `NarrowFloatLoadPolicy`, with `toDTypePolicy()` moved to an extension function. Existing call sites and tests compile unchanged. GGUF: - F16 and BF16 stop unconditionally widening; `Require(BF16)` / `Require(FP16)` are accepted instead of throwing. This is the KEEP_NATIVE GGUF path the CHANGELOG had parked, and it is what makes Require(BF16) real on GGUF. - `keepsNative(policy, native)` is the single decision point. A policy naming one narrow format WIDENS the other rather than keeping it packed. This matters more than it looks: the two encodings are different bit layouts, so mis-tagging F16 bytes as BF16 would decode to silent garbage rather than failing. Enforced by `policy.target == native` in both loaders and covered by a dedicated test on each side. Three pre-existing tests asserted the old limitations (`require_fp16_fails_with_explicit_message`, and GGUF's two `require_*_fails_fast_with_clear_message`). Rewritten to assert the new capability — the behaviour they guarded is gone by design, not by accident. Tests: 4 new SafeTensors FP16 integration tests (packed bytes preserved verbatim, decode bit-identical to the widening path, the two policies independent on a mixed F16+BF16 file) plus 5 policy-mapping tests across both loaders. Existing BF16 suites pass unmodified. apiDump regenerated; full engine build clean (2133 tasks). Refs #884 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Third stage (B2) of #884. FP16 tensors now reach a real kernel instead of paying a per-element decode through `get`, which is what made the B3 KEEP_NATIVE path correct-but-slow for FP16. SPI: - `NarrowFloatMatmulKernel` carries a `codec` and the existing matmul signature; `Bf16MatmulKernel` and the new `Fp16MatmulKernel` extend it. The two formats differ by a decode step and nothing else. - `KernelProvider.matmulFp16()`, defaulting to null so providers without an FP16 kernel cascade as before. Kernels: - `ScalarFp16MatmulKernel` (priority 0, every KMP target) — the correctness reference. - `PanamaVectorFp16MatmulKernel` (priority 50) — vectorized FP32 FMA over n. The DECODE stays scalar into a lane-width scratch buffer: BF16 widens with an integer shift, but binary16 needs exponent rebiasing and mainstream JDKs expose no FP16 vector species. Expect this to trail the BF16 kernel; that is the format on this platform, not a defect here. Dispatch: - `DefaultCpuOpsJvm.chooseQuantizedMatmul` matches `NarrowFloatTensorData` rather than `Bf16TensorData` and picks the kernel by `codec`. Selecting by codec rather than width is load-bearing. Both formats are 2 bytes per element, so a mix-up cannot fail loudly — decoding F16 bytes through the BF16 shift yields plausible-looking wrong numbers with no exception. `fp16_weights_are_not_decoded_by_the_bf16_kernel` pins this from both sides: the dispatch must match the FP16 decode, AND the BF16 decode of the same bytes must differ by >10x, so the test cannot silently pass if the two converge. Accumulation is FP32 in every implementation. The narrow format is a storage width, never an accumulate width. Tests: 11 new (7 parity incl. lane-aligned/tail/k==0/empty shapes and an exact-FP32-reference check, 4 dispatch). Existing BF16 suites pass unmodified (10 parity + 3 dispatch). apiDump regenerated; full engine build clean. NOT included, deliberately: - Native FFM FP16 kernel (priority 100). BF16's lives in native/src/bf16_matmul.c; an FP16 peer needs new C plus build wiring. FP16 tops out at the Panama tier. - Vectorizing the BF16 dequant (IntVector.lanewise(LSHL, 16)) that PanamaVectorBf16MatmulKernel's own KDoc flags — a throughput change to an already-correct kernel, separable from making FP16 work. Refs #884 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ctory Final stage (B4) of #884. Closes the gaps that left BF16 half-supported and repairs a type lattice that disagreed with itself. Lattice — five bugs, four of them code contradicting its own comment: - `FP32.promoteTo(Ternary)` returned Ternary: a promotion to LOWER precision, which no promotion lattice should ever produce. Comment said FP32. - `FP32.isCompatible(Ternary)` returned false; comment said Ternary promotes to FP32. - `FP16.promoteTo(Int32)` returned FP16, silently dropping ~21 bits of integer precision, while `Int32.promoteTo(FP16)` correctly returned FP32. - `Int8.promoteTo(BF16)` returned FP32 even though bf16's 8 significand bits represent every Int8 value exactly — and `BF16.promoteTo(Int8)` said BF16. Now consistent with how Int8 already treats FP16. - `commonPrecisionWith` was a SECOND hand-written lattice with no BF16 arm at all, so `BF16.commonPrecisionWith(Int8)` said FP32 while `BF16.promoteTo(Int8)` said BF16. It now delegates to `promoteTo`, so the two cannot drift apart again. `isConvertibleTo` had no BF16 arm either — every BF16 conversion reported false. Now expressed via a new `DType.isFloatingPoint()` so any float width converts to any other. Completeness: - `supportsFloatOps` admits BF16, which was excluded and therefore left BF16-tagged tensors matmul-only: every elementwise and unary op fell through to the generic scalar path. Callers still re-check for FloatArrayTensorData, so genuinely narrow-storage tensors keep falling through safely. - `convert` to FP16/BF16 now ROUNDS. It previously shared the FP32 branch, so `convert(x, FP16)` only re-tagged the dtype while keeping full FP32 precision — the tensor claimed to be FP16 while holding values no FP16 can represent, with no rounding, no overflow to Inf past the format range, and no NaN handling. BF16 was rejected as a target outright. - `DenseTensorDataFactory`: BF16 arms for `ones`/`full`/`randn`, which threw, AND for both `fromFloatArray` overloads — the latter was not in the original survey and is why the first convert-to-BF16 run still failed. Storage for a narrow `convert` target stays float-backed (values rounded, tag correct), matching how DenseTensorDataFactory already handles narrow tags. Emitting genuinely 2-byte storage from `convert` would change what every existing caller receives and is a separate decision. Tests: 9 lattice-consistency (incl. agreement across all 196 dtype pairs, float/int symmetry, and a "promotion never narrows" invariant) and 8 convert tests (rounding, fp16 overflow to Inf where bf16 stays finite, subnormal underflow, idempotence). apiDump regenerated; full engine build clean. Refs #884 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
📖 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. |
aharakal
approved these changes
Jul 26, 2026
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
2 participants
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.
Implements all four stages of #884. One 16-bit float layer parameterized by codec, rather than two half-built stacks.
Reviewable commit-by-commit — each stage builds clean on its own:
dff4a6b687a28fcedf03f625cb1af5e7Why
Downstream, the FunctionGemma-270M export dequantized Q5_K to FP32, held it through the whole trace, then narrowed straight back to bf16 — with zero arithmetic in between (
VoidTensorOpsis@Backend(id = "void", displayName = "Shape-only")). Measured: 324 weight globals, 436,111,680 elements = 1.62 GiB FP32 against an 832 MiB bf16 destination, ~4.3 GiB peak live. A ×6.5 inflation to squeeze straight back down, because nothing downstream could carry narrow dense storage.What landed
B1 —
NarrowFloatCodecwithBf16Codec/Fp16Codec(pure integer bit math incommonMain, so every KMP target gets it);NarrowFloatDenseTensorData;Fp16DenseTensorData.B3 — SafeTensors gains the F16
KEEP_NATIVEarm BF16 already had; GGUF stops unconditionally widening F16/BF16.Require(FP16)andRequire(BF16)are accepted instead of throwing.B2 —
NarrowFloatMatmulKernelSPI carrying a codec;ScalarFp16MatmulKernel+PanamaVectorFp16MatmulKernel; dispatch widened fromis Bf16TensorDatatois NarrowFloatTensorData, selecting by codec.B4 —
supportsFloatOpsadmits BF16;convertrounds properly and accepts BF16;DenseTensorDataFactoryBF16 arms; the dtype lattice repaired.Result: an F16/BF16 checkpoint loads with
Require(...), stays at 2 bytes per element, dispatches to a codec-selected kernel, and traces with the correct dtype into MLIR — no widening anywhere.Eight latent bugs fixed along the way
Not the assignment, but found while wiring this up:
TraceSessionhad noBF16arm, so BF16 tensors traced as FP32. This is why the downstream exporter has to regex-rewritexf32>→xbf16>in emitted MLIR text — it unblocks SKaiNET-transformers#248.TensorStorageFactorydecoded 2-byte narrow data with a hard 4-byte reader, yielding half the element count with each "float" assembled from two adjacent elements. Readers already taggedDense(2); only the decode ignored it. Silent garbage, not an error.FP32.promoteTo(Ternary)returnedTernary— a promotion to lower precision. Its own comment said FP32.FP32.isCompatible(Ternary)returnedfalsewhile its comment said Ternary promotes to FP32.FP16.promoteTo(Int32)returnedFP16, dropping ~21 bits, whileInt32.promoteTo(FP16)correctly said FP32.Int8.promoteTo(BF16)returnedFP32though bf16's 8 significand bits hold every Int8 value exactly — andBF16.promoteTo(Int8)said BF16.commonPrecisionWithwas a second lattice with no BF16 arm, disagreeing withpromoteTo. Now delegates, so they cannot drift.isConvertibleTohad no BF16 arm — every BF16 conversion reported false.Review notes
Bf16Codec.encodetruncates rather than rounding. Deliberate and contractual: it matches the existingfloatToBf16Bits, the safetensors writer, and the on-device A/B that qualified bf16 as a bit-exact drop-in for the f16 vmfb. A regression test pins the two implementations together so a future "fix" to rounding fails loudly.Kernel selection is by codec, not width. Both formats are 2 bytes per element, so a mix-up cannot fail loudly — decoding F16 bytes through the BF16 shift gives plausible-looking wrong numbers.
fp16_weights_are_not_decoded_by_the_bf16_kernelpins it from both sides: the dispatch must match the FP16 decode, and the BF16 decode of the same bytes must differ by >10×, so the test cannot silently pass if they converge. The same applies in the loaders: a policy naming one narrow format widens the other rather than mis-tagging it.Accumulation is FP32 everywhere. The narrow format is a storage width, never an accumulate width.
convertto a narrow target keeps float-backed storage (values rounded, tag correct), matching howDenseTensorDataFactoryalready handles narrow tags. Emitting genuinely 2-byte storage fromconvertwould change what every existing caller receives — a separate decision.Three pre-existing tests asserted the old limitations (
require_fp16_fails_with_explicit_message, GGUF's tworequire_*_fails_fast_with_clear_message). Rewritten to assert the new capability; the behaviour they guarded is gone by design.Verification
78 new tests. Full engine build clean at 2133 tasks after every stage;
apiDumpregenerated andapiCheckgreen; every pre-existing BF16 suite (Bf16TensorDataTest,PanamaVectorBf16MatmulKernelParityTest,NativeBf16MatmulKernelParityTest,Bf16MatmulDispatchTest, the SafeTensors BF16 policy tests) passes unmodified.Two exclusions needed for a green local
build, both failing identically on cleandevelopand unrelated to this work::skainet-apps:skainet-grayscale-cli:test(testGpuCapabilitiesDetectionasserts CUDA/IREE are absent; this machine has them) andjsBrowserTest(no ChromeHeadless installed).Not included
IntVector.lanewise(LSHL, 16)) thatPanamaVectorBf16MatmulKernel's KDoc flags — throughput only, on an already-correct kernel.Closes #884
🤖 Generated with Claude Code