transpose() has no NarrowFloatTensorData arm — widens FP16/BF16 weights elementwise, 0.2-4.5s per projection

[michalharakal]

Summary

DefaultCpuOps.transpose has a lazy-transpose arm for every packed quant type, but none for NarrowFloatTensorData. FP16/BF16 weights therefore fall through to the generic transpose, which walks the tensor element by element through boxed get(), decoding each 2-byte element individually.

Because both Linear.onForward and (downstream) LlamaRuntime.linearProject transpose the weight on every call, this happens per weight, per token. At real model sizes KEEP_NATIVE is not merely un-accelerated — it is unusably slow.

Follow-up to #884 / #886. Related to #887, but independent: #887 is about the FP16 decode being slow, this is about the narrow kernel never being reached at all.

Measurements

Intel i7-9750H (AVX2), OpenJDK 21.0.11, engine at ddad83a. x.matmul(w.t()) with an FP16 weight in the real [out, in] orientation, versus the same matmul with the weight already [in, out] so no transpose is needed:

shape (in x out)	batch	matmul only, [in, out]	.t() + matmul, [out, in]
2048 x 2048	1	18.6 ms	206 ms
4096 x 4096	1	73.8 ms	1 319 ms
4096 x 11008	1	197.3 ms	4 367 ms
11008 x 4096	1	199.0 ms	2 196 ms

The transpose is the bulk of the difference. Timings are nearly flat across batch size, confirming the cost is the one-off weight transpose rather than the matmul.

Root cause

skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt, in transpose:

when (val d = tensor.data) {
    is Q4_KTensorData -> return newTensor(Q4_KBlockTensorData(Shape(cols, rows), d.packedData) ...)
    is Q5_KTensorData -> ...
    is Q6_KTensorData -> ...
    is Q5_1TensorData -> ...
    is Q5_0TensorData -> ...
    is Q8_0TensorData -> ...
    is Q4_0TensorData -> ...
    else -> {}   // <- NarrowFloatTensorData lands here
}

The comment on that block says it "now covers every quant type chooseQuantizedMatmulHeap dispatches — i.e. every packed type that can be a matmul weight". Narrow floats became such a type in #886 and were not added.

DefaultCpuOpsJvm.transpose intercepts separately (Q4/Q8 MemorySegment) before delegating, so it needs the same arm.

Why a pure shape swap is not enough here

For the K-quants the lazy transpose is valid because the bytes were already relaid into input-block-major order at load (relayoutKSeriesRowMajorToBlockMajor on the consumer side), so flipping the shape lands on the layout the kernel indexes.

Narrow-float bytes are plain dense row-major. Reinterpreting a row-major [out, in] buffer as [in, out] yields a different matrix, not its transpose — it would produce silently wrong results rather than an error. So there are two viable shapes for the fix:

1. Lazy shape swap in the engine + byte relayout at load on the consumer side. Mirrors the K-quant design exactly; transpose stays free. Requires consumers to store narrow weights [in, out] while declaring [out, in], so it is a coordinated change.
2. Real 2-byte element transpose in the engine. Self-contained, no consumer change, allocates a new packed buffer per call — but an O(n) System.arraycopy-class permutation over 2-byte elements is on the order of 100x cheaper than the current boxed elementwise path, so even this would remove most of the regression.

Option 1 is the better ceiling; option 2 is the safer immediate fix and could land first. Whichever is chosen, codec must be carried across — NarrowFloatDenseTensorData(Shape(cols, rows), d.packedData, d.codec) — since FP16 and BF16 are both 2 bytes per element and swapping them produces plausible wrong numbers rather than an exception.

Suggested interim mitigation

Until this lands, consumers enabling DTypePolicy.Require/Prefer on a narrow format get a large silent slowdown at model scale. A note in the NarrowFloatLoadPolicy / SafeTensorsParametersLoader KDoc that KEEP_NATIVE currently saves at-rest memory at a significant per-token cost would prevent surprises.

Reproducing

NarrowFloatMatmulBenchmark in SKaiNET-transformers (llm-inference/llama/src/jvmTest, branch chore/validate-engine-0.38.0-snapshot), transpose column:

./gradlew :llm-inference:llama:jvmTest --tests "*NarrowFloatMatmulBenchmark*" \
  -PuseMavenLocalSkainet=true -Dskainet.bench.narrow=true --rerun-tasks -i