Skip to content

BF16 dispatch chain (Phase 1/3): add Bf16TensorData + Bf16DenseTensorData - #610

Merged
michalharakal merged 1 commit into
developfrom
feature/bf16-tensor-data
May 16, 2026
Merged

BF16 dispatch chain (Phase 1/3): add Bf16TensorData + Bf16DenseTensorData#610
michalharakal merged 1 commit into
developfrom
feature/bf16-tensor-data

Conversation

@michalharakal

@michalharakal michalharakal commented May 16, 2026

Copy link
Copy Markdown
Contributor

Resolves #609. First of three follow-ups to #605 (BF16 matmul kernels).

Why

The BF16 matmul kernels landed in #605 (scalar / Panama / native) are plumbed but can't be reached from ops.matmul today — the SafeTensors loader unconditionally dequants BF16 weights to FP32 at load, so by the time DefaultCpuOps.matmul sees the operand it's a FloatArrayTensorData rather than anything that the new Bf16MatmulKernel recognises.

This PR adds the missing TensorData type so the rest of the chain has something to dispatch on. No loader or dispatch changes here — those are Phase 2 and Phase 3.

What

Bf16TensorData : TensorData<DType, Float> plus a concrete Bf16DenseTensorData(shape, ByteArray) in commonMain. Surface area:

API purpose
Bf16TensorData.packedData: ByteArray zero-copy access to packed 2-bytes-per-element BF16; what the matmul kernel reads
get(*indices): Float decode-on-read; makes the Tensor surface look like FP32 for callers that don't need the raw bytes
set(*indices, value: Float) truncate FP32 → BF16 (lossy by construction)
copyToFloatArray() bulk decode
Bf16TensorData.floatToBf16Bits / bf16BitsToFloat (companion) bit-shift helpers — same identity used by bf16_matmul.c and PanamaVectorBf16MatmulKernel
Bf16DenseTensorData.fromFloatArray(shape, FloatArray) constructor for tests + offline round-trips

Mirrors the structure of Q8_0BlockTensorData minus the block accessors (BF16 is dense, no per-block scale).

Tests

11 unit tests in Bf16TensorDataTest (commonTest):

  • Round-trip set/get within BF16 precision (1e-2 absolute).
  • Raw byte-order check: FP32 1.0 = BF16 0x3F80 = bytes [0x80, 0x3F].
  • Signed-zero preservation (BF16 0x8000 vs 0x0000).
  • 2D and 3D shape × stride correctness.
  • Bulk copyToFloatArray() parity with element-by-element get.
  • Undersized buffer / out-of-bounds / wrong-rank rejections.
  • Bit-identity for FP32 values whose low 16 bits are zero (e.g. 0.0, 1.0, 0.5, 256.0).

All pass on jvmTest and linuxX64Test. (The unrelated kotlinWasmStoreYarnLock failure is pre-existing on develop and not touched by this PR.)

Phase 2 / 3 (follow-up PRs)

phase what
2 SafeTensors loader opt-in to skip dequant for BF16 tensors and emit Bf16DenseTensorData instead. Adds a precision-policy knob so existing consumers stay unaffected.
3 DefaultCpuOpsJvm.chooseQuantizedMatmul dispatch for Bf16TensorData via the Bf16MatmulKernel SPI. Mirrors the Q8_0 dispatch wiring (#608).
4 (optional) Gemma-3n end-to-end smoke. Probably lives in SKaiNET-transformers.

🤖 Generated with Claude Code

@michalharakal @claude
Add Bf16TensorData + Bf16DenseTensorData (#609)
d974ed9 
Phase 1 of the three-phase BF16 dispatch chain. Foundation only — no
loader or dispatch changes here; those follow in Phase 2 / Phase 3.

`Bf16TensorData : TensorData<DType, Float>` is the recognition surface
for "this weight is packed BF16 bytes" so the upcoming dispatch in
`DefaultCpuOpsJvm.chooseQuantizedMatmul` (Phase 3) can route via the
`Bf16MatmulKernel` SPI without dequant-to-FP32 at load.

Surface area:

  - `Bf16TensorData` interface (commonMain) — exposes `packedData:
    ByteArray` (2 bytes per element, little-endian) for zero-copy
    hand-off to SIMD matmul kernels, plus `Bf16TensorData.Companion
    .floatToBf16Bits` / `bf16BitsToFloat` static helpers.
  - `Bf16DenseTensorData` concrete impl — backed by a packed ByteArray.
    `get(*indices): Float` decodes BF16 → FP32 on read; `set` truncates
    FP32 → BF16 (lossy by construction). Bulk `copyToFloatArray()` for
    consumers that just want all values.
  - `Bf16DenseTensorData.fromFloatArray(shape, FloatArray)` factory for
    tests and offline round-tripping.
  - `Bf16TensorData.toFloatArray()` extension for dequant fallback.

11 unit tests in commonTest cover: round-trip within BF16 precision
(1e-2 abs), raw byte-order check (FP32 1.0 = BF16 0x3F80 = bytes
[0x80, 0x3F]), get/set primitives, signed zero preservation, 2D /
3D shape strides, bulk copyToFloatArray parity with element-by-
element decode, undersized-buffer / out-of-bounds / wrong-rank
rejections, and bit-identity for FP32 values that have zero in the
low 16 bits.

Passes jvmTest + linuxX64Test on the lang-core module. Refs #609.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@github-actions

github-actions Bot commented May 16, 2026

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-610 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

@michalharakal
michalharakal merged commit d23041c into develop May 16, 2026
10 of 11 checks passed
@michalharakal
michalharakal deleted the feature/bf16-tensor-data branch May 16, 2026 21:12
This was referenced May 16, 2026
Closed
Merged
Closed
Merged
MacOS pushed a commit to MacOS/SKaiNET that referenced this pull request Jul 10, 2026
@michalharakal @claude
Add Bf16LoadPolicy.KEEP_NATIVE to SafeTensorsParametersLoader (SKaiNE…
cecac70 
…T-developers#611)

Phase 2 of the three-phase BF16 dispatch chain. Follow-up to SKaiNET-developers#610
(Bf16TensorData merged).

Today's loader unconditionally dequants BF16 → FP32 at load
(`SafeTensorsParametersLoader.kt` line 87, `dequantBF16(bytes)`),
which means even with the new Bf16TensorData type in place no
SafeTensors-loaded weight ever reaches it. This PR adds an opt-in
policy so consumers that want native BF16 (Gemma-3n is the obvious
first one) get a Bf16DenseTensorData-backed tensor; everyone else
stays on the dequant path with zero behavioural change.

Surface:
  - new `Bf16LoadPolicy` enum (commonMain) with `DEQUANT_TO_FP32`
    (default) and `KEEP_NATIVE` cases. Documents the trade-off:
    memory halved vs. per-element decode cost on non-matmul ops.
  - new constructor parameter `bf16Policy: Bf16LoadPolicy =
    Bf16LoadPolicy.DEQUANT_TO_FP32`. Default preserves source +
    bytecode compat for every existing Kotlin caller.
  - new branch in the `DataType.BFLOAT16` case: when policy is
    `KEEP_NATIVE`, wrap the on-disk bytes in `Bf16DenseTensorData`
    and emit via `ctx.fromData(...)`. The consumer-visible dtype
    stays `FP32::class` (same pattern as Q4_K / Q8_0 tensors —
    quantised storage, FP32 dtype tag); only `tensor.data` differs.

4 new tests in `commonTest`:
  - DEQUANT_TO_FP32 path produces FloatArrayTensorData with values
    within BF16 precision.
  - KEEP_NATIVE path produces Bf16DenseTensorData whose packedData
    byte array matches the on-disk bytes verbatim.
  - Decoded values from both paths are bit-identical (both apply
    the same `bf16_bits << 16` math; only WHEN differs).
  - Mixed BF16+FP32 file under KEEP_NATIVE — BF16 becomes
    Bf16DenseTensorData, FP32 stays FloatArrayTensorData.

Refs SKaiNET-developers#611. Full `:skainet-io:skainet-io-safetensors:jvmTest` suite
passes on linux-x86_64 / JDK 21.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
MacOS pushed a commit to MacOS/SKaiNET that referenced this pull request Jul 10, 2026
@michalharakal @claude
Wire BF16 matmul dispatch in DefaultCpuOpsJvm.chooseQuantizedMatmul (S…
c4de7fb 
…KaiNET-developers#613)

Final phase of the three-phase BF16 dispatch chain. Follow-ups to
SKaiNET-developers#610 (Bf16TensorData) and SKaiNET-developers#612 (loader KEEP_NATIVE policy) — both
merged.

After this PR, a consumer that flips `bf16Policy = KEEP_NATIVE` on
SafeTensorsParametersLoader (or constructs a `Bf16DenseTensorData`
directly) gets the SIMD-vectorised BF16 matmul path with zero other
code changes. Native FFM kernel (priority 100) wins when the bundled
libskainet_kernels.so is loaded; falls through to Panama Vector (50)
then to the scalar SPI reference (0).

Implementation:
  - New `bf16MatmulKernel: Bf16MatmulKernel` lazy in DefaultCpuOpsJvm.
    Non-null with `ScalarBf16MatmulKernel` floor — mirrors
    `fp32MatmulKernel`'s pattern rather than the nullable
    `q4kMatmulKernel` / `q8_0MatmulKernel` pattern (which exist because
    Q4_K / Q8_0 have legacy non-SPI fallbacks via
    `JvmQuantizedVectorKernels`; BF16 has no such legacy).
  - New `is Bf16TensorData ->` branch in `chooseQuantizedMatmul`'s
    `when (bData)` block. The BF16 SPI kernel is a full SGEMM
    `(m, n, k)` with byte-strides on the B operand — no per-batch
    matvec loop like Q4_K/Q8_0/Q6_K need.

3 integration tests in `Bf16MatmulDispatchTest`:
  - single-batch matmul (`[1, k] × [k, n]` BF16) matches scalar
    reference within `1e-2 * k`.
  - multi-batch matmul (`m=3, k=256, n=32`) — exercises a 2D output.
  - LLM-typical 512² attention projection.

Refs SKaiNET-developers#613. Full `:skainet-backends:skainet-backend-cpu:jvmTest` and
`:skainet-backends:skainet-backend-native-cpu:jvmTest` suites pass on
linux-x86_64 / JDK 21 with `--add-modules jdk.incubator.vector`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

BF16 dispatch chain (Phase 1/3): add Bf16TensorData + Bf16DenseTensorData

1 participant

@michalharakal