Skip to content

Packed Q5_1 / Q5_0 SIMD matmul kernel + lazy transpose (CPU backend) #708

Description

@michalharakal

Summary

Add a packed SIMD matmul path for Q5_1 (and Q5_0 for completeness) in the CPU backend, mirroring the existing Q4_K / Q6_K chain, so quantized weights of these types can be consumed packed (FP32 activations × packed quant weights) instead of being dequantized to FP32.

Motivation

Downstream in SKaiNET-transformers, loading a real Gemma GGUF whose attention/FFN weights are Q5_1 (e.g. functiongemma-270m, marketed "Q5_K_M") under QuantPolicy.NATIVE_OPTIMIZED had no packed kernel for Q5_1. The transformers-side converter therefore either left raw bytes (crashing linearProject -> ops.transpose with "Transpose requires at least 2 dimensions") or, after the fix below, dequantizes Q5_1 to FP32 — correct, but it forgoes the packed memory/speed benefit the Q4_K/Q6_K/Q4_0/Q8_0 tensors keep.

Measured on functiongemma-270m: ~0.67 tok/s with the FP32 dequant fallback vs ~0.23 when forcing a global dequant; a packed Q5_1 kernel should close the remaining gap to the fully-packed models.

Scope (CPU backend)

Mirror the Q4_K / Q6_K implementation:

  1. Packed tensor-data type for Q5_1 (cf. Q4_KBlockTensorData / Q6_KBlockTensorData in sk.ainet.lang.tensor.data), carrying the packed bytes + logical 2-D shape.
  2. matmulQ5_1Vec in skainet-backends/skainet-backend-cpu/.../ops/JvmQuantizedVectorKernels.kt — FP32 activations × packed Q5_1 weights, alongside matmulQ4_KVec / matmulQ6_KVec.
  3. Lazy transpose branch for the new type in DefaultCpuOpsJvm.transpose (shape-swap, no data copy), matching the existing Q4_K / Q6_K / Q4_0 / Q8_0 branches, so linearProject (x @ W.t()) flows through without a FP32 round-trip.
  4. Wire matmul dispatch (DefaultCpuOpsJvm.matmul) to pick the Q5_1 kernel when it sees the packed data type.

(Once this lands, the transformers converter GemmaMemSegConverter.convertOne gets a Q5_1/Q5_0 packed case that replaces the dequant fallback — tracked in transformers#170.)

Acceptance

  • A kernel unit test: FP32 activations × packed Q5_1 weights matches an FP32-dequant reference to the same tolerance as the existing matmulQ4_KVec / matmulQ6_KVec tests.
  • ops.transpose on the packed Q5_1 type is a lazy shape-swap (no copy), verified by a small backend test.
  • End-to-end (via transformers#170): a Q5_1 Gemma checkpoint runs under NATIVE_OPTIMIZED with those weights packed, output matching the DEQUANTIZE_TO_FP32 reference token-for-token.

References (CPU backend)

  • skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmQuantizedVectorKernels.ktmatmulQ4_KVec / matmulQ6_KVec
  • skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt — lazy transpose branches (Q4_KTensorData / Q6_KTensorData / Q4MemorySegmentMarker / Q8MemorySegmentMarker)
  • sk.ainet.lang.tensor.data.Q4_KBlockTensorData / Q6_KBlockTensorData — packed tensor-data pattern to mirror

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions