perf(cuda): separated GPU layout for Q5_0 GEMV

Q5_0 is the dominant weight type in Gemma3-1B Q4_K_M (117 of 170
weight tensors). The interleaved 22-byte block layout required
byte-wise __ldg loads on ARM64 Grace Hopper (blocks not 4-byte
aligned after block 0). This caused ~40% throughput regression.

Introduce a separated GPU layout (scales | qh | qs) where each
region is naturally aligned. The GEMV kernel now reads fp16 scales
at 2-byte boundaries and uint32 qh at 4-byte boundaries with single
__ldg instructions instead of 4-byte-at-a-time reconstruction.

Also add RawBytesGPU() to Q5_0Storage for the separated layout,
matching the pattern used by Q4Storage.

Author: dndungu

compute/gpu_engine.go

@@ -393,13 +393,14 @@ func (e *GPUEngine[T]) UploadWeights(tensors []*tensor.TensorNumeric[float32]) e
 			q4Uploaded++
 			continue
 		}
-		// Upload Q5_0 raw bytes to GPU for fused GEMV kernel.
-		// Q5_0 blocks (22 bytes per 32 values) are uploaded contiguously.
+		// Upload Q5_0 in separated GPU layout (scales | qh | qs) for fast GEMV.
+		// The separated layout aligns fp16 and uint32 fields naturally, avoiding
+		// byte-wise loads on ARM64 Grace Hopper.
 		if qs, ok := any(t.GetStorage()).(*tensor.Q5_0Storage); ok {
 			if ptr, _, _ := qs.GPUPtr(); ptr != nil {
 				continue // already on GPU
 			}
-			rawBytes := qs.RawBytes()
+			rawBytes := qs.RawBytesGPU()
 			devPtr, err := e.allocWeight(len(rawBytes))
 			if err != nil {
 				return fmt.Errorf("alloc Q5_0 GPU (shape %v): %w", t.Shape(), err)
@@ -2133,11 +2134,12 @@ func (e *GPUEngine[T]) matMulQ5_0(ctx context.Context, qs *tensor.Q5_0Storage, a
 
 	var devW unsafe.Pointer
 	var freeW func()
+	nBlocks := qs.NumBlocks()
 	if ptr, _, _ := qs.GPUPtr(); ptr != nil {
 		devW = ptr
 		freeW = func() {}
 	} else {
-		rawBytes := qs.RawBytes()
+		rawBytes := qs.RawBytesGPU()
 		var err error
 		devW, err = e.pool.Alloc(e.deviceID, len(rawBytes))
 		if err != nil {
@@ -2151,6 +2153,9 @@ func (e *GPUEngine[T]) matMulQ5_0(ctx context.Context, qs *tensor.Q5_0Storage, a
 	}
 	defer freeW()
 
+	qhOff := tensor.Q5_0GPUQhOffset(nBlocks)
+	qsOff := tensor.Q5_0GPUQsOffset(nBlocks)
+
 	if n == 1 {
 		devX, cleanupX, err := getDevicePtr(e, b)
 		if err != nil {
@@ -2165,7 +2170,7 @@ func (e *GPUEngine[T]) matMulQ5_0(ctx context.Context, qs *tensor.Q5_0Storage, a
 			return e.cpu.MatMul(ctx, a, b, dst...)
 		}
 
-		if err := e.kernels.GemvQ5_0F32(devW, devX, devY, m, k, e.stream); err != nil {
+		if err := e.kernels.GemvQ5_0F32(devW, devX, devY, m, k, qhOff, qsOff, e.stream); err != nil {
 			e.pool.Free(e.deviceID, devY, cSize)
 			return e.cpu.MatMul(ctx, a, b, dst...)
 		}
@@ -2240,11 +2245,12 @@ func (e *GPUEngine[T]) matMulQ5_0BWeight(ctx context.Context, a *tensor.TensorNu
 
 	var devQ5_0 unsafe.Pointer
 	var freeQ5_0 func()
+	nBlocks := qs.NumBlocks()
 	if ptr, _, _ := qs.GPUPtr(); ptr != nil {
 		devQ5_0 = ptr
 		freeQ5_0 = func() {}
 	} else {
-		rawBytes := qs.RawBytes()
+		rawBytes := qs.RawBytesGPU()
 		var err error
 		devQ5_0, err = e.pool.Alloc(e.deviceID, len(rawBytes))
 		if err != nil {
@@ -2258,6 +2264,9 @@ func (e *GPUEngine[T]) matMulQ5_0BWeight(ctx context.Context, a *tensor.TensorNu
 	}
 	defer freeQ5_0()
 
+	qhOff := tensor.Q5_0GPUQhOffset(nBlocks)
+	qsOff := tensor.Q5_0GPUQsOffset(nBlocks)
+
 	if m == 1 {
 		devX, cleanupX, err := getDevicePtr(e, a)
 		if err != nil {
@@ -2272,7 +2281,7 @@ func (e *GPUEngine[T]) matMulQ5_0BWeight(ctx context.Context, a *tensor.TensorNu
 			return e.cpu.MatMul(ctx, a, b, dst...)
 		}
 
-		if err := e.kernels.GemvQ5_0F32(devQ5_0, devX, devY, n, k, e.stream); err != nil {
+		if err := e.kernels.GemvQ5_0F32(devQ5_0, devX, devY, n, k, qhOff, qsOff, e.stream); err != nil {
 			e.pool.Free(e.deviceID, devY, cSize)
 			return e.cpu.MatMul(ctx, a, b, dst...)
 		}

internal/cuda/kernels/gemv_q5_0.cu

@@ -1,13 +1,13 @@
 /* Q5_0 fused dequant-GEMV kernel for single-token decode (batch=1).
  *
- * Reads Q5_0 blocks directly, dequantizes in registers (no global
- * memory intermediary), multiplies by the activation vector, and accumulates
- * in FP32. This halves memory traffic compared to separate dequant + GEMV.
+ * GPU-optimized SEPARATED layout (from Q5_0Storage.RawBytesGPU):
+ *   Region 1: [nBlocks * 2 bytes] fp16 scales, padded to 16-byte boundary
+ *   Region 2: [nBlocks * 4 bytes] uint32 qh values, padded to 16-byte boundary
+ *   Region 3: [nBlocks * 16 bytes] packed nibbles (qs)
  *
- * Q5_0 block (22 bytes, 32 values):
- *   [0:2]   fp16 d      -- block scale
- *   [2:6]   uint32 qh   -- 32 high bits (one per element)
- *   [6:22]  16 bytes qs  -- packed nibbles (two 4-bit values per byte)
+ * This layout ensures natural alignment: fp16 at 2-byte, uint32 at 4-byte.
+ * Eliminates the byte-wise loads required for the interleaved 22-byte layout
+ * on ARM64 Grace Hopper.
  *
  * Dequantization (matching llama.cpp dequantize_row_q5_0):
  *   For j in 0..15:
@@ -25,26 +25,22 @@
 #include <stdint.h>
 
 #define Q5_0_BLOCK_SIZE  32
-#define Q5_0_BLOCK_BYTES 22
 #define Q5_0_WARPS_PER_BLOCK 4
 #define Q5_0_WARP_SIZE       32
 
-/* ---------- Fused GEMV kernel ----------
+/* ---------- Fused GEMV kernel (separated GPU layout) ----------
  *
  * y[row] = sum_k dequant(W_q5_0[row, k]) * x[k]
  *
- * Strategy:
- *   - Load input vector x into shared memory (all threads cooperate).
- *   - One warp per row for simplicity and good occupancy.
- *   - Each lane processes a strided subset of blocks.
- *   - Within each block, 16 packed bytes yield 32 dequantized values.
- *   - Warp shuffle reduction produces the final dot product.
+ * W_q5_0 points to the separated layout base. qhOffset and qsOffset
+ * are byte offsets to the qh and qs regions respectively.
  */
 __global__ void gemv_q5_0_kernel(
     const uint8_t* __restrict__ W_q5_0,
     const float*   __restrict__ x,
     float*         __restrict__ y,
-    int M, int K)
+    int M, int K,
+    int qhOffset, int qsOffset)
 {
     extern __shared__ float sx[];
 
@@ -62,27 +58,20 @@ __global__ void gemv_q5_0_kernel(
     if (row >= M) return;
 
     int blocks_per_row = K / Q5_0_BLOCK_SIZE;
-    const uint8_t* row_data = W_q5_0 + (size_t)row * blocks_per_row * Q5_0_BLOCK_BYTES;
+
+    /* Pointers to the three separated regions for this row. */
+    const __half*    row_scales = (const __half*)(W_q5_0 + row * blocks_per_row * 2);
+    const uint32_t*  row_qh    = (const uint32_t*)(W_q5_0 + qhOffset + row * blocks_per_row * 4);
+    const uint8_t*   row_qs    = W_q5_0 + qsOffset + (size_t)row * blocks_per_row * 16;
 
     float acc = 0.0f;
 
     /* Each lane handles a strided subset of blocks. */
     for (int bi = lane_id; bi < blocks_per_row; bi += Q5_0_WARP_SIZE) {
-        const uint8_t* blk = row_data + bi * Q5_0_BLOCK_BYTES;
-
-        /* Read fp16 d using byte-wise load (ARM64 alignment safety).
-         * Q5_0 blocks are 22 bytes — not a multiple of 4, so blk may
-         * be misaligned for uint16/uint32 casts after the first block. */
-        uint16_t d_bits = (uint16_t)__ldg(&blk[0]) | ((uint16_t)__ldg(&blk[1]) << 8);
-        float d = __half2float(*reinterpret_cast<const __half*>(&d_bits));
-
-        /* Read qh (32 high bits) using byte-wise load. */
-        uint32_t qh = (uint32_t)__ldg(&blk[2])
-                     | ((uint32_t)__ldg(&blk[3]) << 8)
-                     | ((uint32_t)__ldg(&blk[4]) << 16)
-                     | ((uint32_t)__ldg(&blk[5]) << 24);
-
-        const uint8_t* qs = blk + 6;
+        /* All loads are naturally aligned in the separated layout. */
+        float d = __half2float(__ldg(&row_scales[bi]));
+        uint32_t qh = __ldg(&row_qh[bi]);
+        const uint8_t* qs = row_qs + bi * 16;
         int k_base = bi * Q5_0_BLOCK_SIZE;
 
         /* Process 16 packed bytes -> 32 dequantized values. */
@@ -119,6 +108,7 @@ __global__ void gemv_q5_0_kernel(
 extern "C" cudaError_t gemv_q5_0_f32(
     const void* W_q5_0, const float* x, float* y,
     int M, int K,
+    int qhOffset, int qsOffset,
     cudaStream_t stream)
 {
     if (K % Q5_0_BLOCK_SIZE != 0) {
@@ -130,7 +120,7 @@ extern "C" cudaError_t gemv_q5_0_f32(
     int smem = K * sizeof(float);
 
     gemv_q5_0_kernel<<<grid, threads, smem, stream>>>(
-        (const uint8_t*)W_q5_0, x, y, M, K);
+        (const uint8_t*)W_q5_0, x, y, M, K, qhOffset, qsOffset);
 
     return cudaGetLastError();
 }

internal/cuda/kernels/gemv_q5_0.go

@@ -14,17 +14,17 @@ import (
 )
 
 // GemvQ5_0F32 performs Q5_0 fused dequant-GEMV: y = dequant(W_q5_0) * x.
-// W_q5_0 is raw Q5_0 blocks for matrix [M, K] (row-major block layout).
-// x is [K] FP32 input vector. y is [M] FP32 output vector.
-// K must be a multiple of 32.
+// W_q5_0 is the separated GPU layout (scales | qh | qs).
+// qhOffset and qsOffset are byte offsets to the qh and qs regions.
 func GemvQ5_0F32(
 	W_q5_0, x, y unsafe.Pointer,
-	M, K int,
+	M, K, qhOffset, qsOffset int,
 	stream unsafe.Pointer,
 ) error {
 	err := C.gemv_q5_0_f32(
 		W_q5_0, (*C.float)(x), (*C.float)(y),
 		C.int(M), C.int(K),
+		C.int(qhOffset), C.int(qsOffset),
 		C.cudaStream_t(stream),
 	)
 	if err != C.cudaSuccess {

internal/cuda/kernels/gemv_q5_0.h

@@ -1,13 +1,11 @@
-/* Q5_0 fused dequant-GEMV kernel interface.
+/* Q5_0 fused dequant-GEMV kernel interface (separated GPU layout).
  *
- * Q5_0 block format (22 bytes per 32 values):
- *   - 2 bytes: fp16 d (block scale)
- *   - 4 bytes: uint32 qh (32 high bits, one per element)
- *   - 16 bytes: qs (packed nibbles, two 4-bit values per byte)
+ * GPU layout (from Q5_0Storage.RawBytesGPU):
+ *   Region 1: [nBlocks * 2 bytes] fp16 scales, padded to 16-byte boundary
+ *   Region 2: [nBlocks * 4 bytes] uint32 qh values, padded to 16-byte boundary
+ *   Region 3: [nBlocks * 16 bytes] packed nibbles (qs)
  *
  * Computes: y[m] = sum_k( dequant(W_q5_0[m,k]) * x[k] )
- * W_q5_0 is raw Q5_0 blocks laid out row-major.
- * x is [K] FP32 input vector. y is [M] FP32 output vector.
  * Batch=1 only (GEMV, not GEMM).
  */
 #ifndef GEMV_Q5_0_H
@@ -22,16 +20,18 @@ extern "C" {
 /* gemv_q5_0_f32 performs Q5_0 fused dequant-GEMV:
  *   y[m] = sum_k( dequant(W_q5_0[m,k]) * x[k] )
  *
- * W_q5_0: device pointer to raw Q5_0 blocks for matrix W [M, K].
- *         M * ceil(K/32) blocks, each 22 bytes. Row-major layout.
- * x:      device pointer to [K] float input vector.
- * y:      device pointer to [M] float output vector.
- * M, K:   matrix dimensions. K must be a multiple of 32.
- * stream: CUDA stream.
+ * W_q5_0:   device pointer to separated Q5_0 layout (scales | qh | qs).
+ * x:        device pointer to [K] float input vector.
+ * y:        device pointer to [M] float output vector.
+ * M, K:     matrix dimensions. K must be a multiple of 32.
+ * qhOffset: byte offset from W_q5_0 to the qh region.
+ * qsOffset: byte offset from W_q5_0 to the qs region.
+ * stream:   CUDA stream.
  */
 cudaError_t gemv_q5_0_f32(
     const void* W_q5_0, const float* x, float* y,
     int M, int K,
+    int qhOffset, int qsOffset,
     cudaStream_t stream);
 
 #ifdef __cplusplus

internal/cuda/kernels/gemv_q5_0_purego.go

@@ -10,10 +10,11 @@ import (
 )
 
 // GemvQ5_0F32 performs Q5_0 fused dequant-GEMV: y = dequant(W_q5_0) * x.
-// W_q5_0 is raw Q5_0 blocks, x is [K] FP32, y is [M] FP32.
+// W_q5_0 is the separated GPU layout (scales | qh | qs).
+// qhOffset and qsOffset are byte offsets to the qh and qs regions.
 func GemvQ5_0F32(
 	W_q5_0, x, y unsafe.Pointer, //nolint:gocritic // match CGo API
-	M, K int, //nolint:gocritic // match CGo API
+	M, K, qhOffset, qsOffset int, //nolint:gocritic // match CGo API
 	stream unsafe.Pointer,
 ) error {
 	k := klib()
@@ -22,6 +23,8 @@ func GemvQ5_0F32(
 	}
 	ret := cuda.Ccall(k.launchGemvQ5_0F32,
 		uintptr(W_q5_0), uintptr(x), uintptr(y),
-		uintptr(M), uintptr(K), uintptr(stream))
+		uintptr(M), uintptr(K),
+		uintptr(qhOffset), uintptr(qsOffset),
+		uintptr(stream))
 	return checkKernel(ret, "gemv_q5_0_f32")
 }

internal/gpuapi/cuda_kernels.go

@@ -129,8 +129,8 @@ func (k *CUDAKernels) GemvQ6KF32(wQ6K, x, y unsafe.Pointer, M, K int, s Stream)
 	return kernels.GemvQ6KF32(wQ6K, x, y, M, K, streamPtr(s))
 }
 
-func (k *CUDAKernels) GemvQ5_0F32(wQ5_0, x, y unsafe.Pointer, M, K int, s Stream) error {
-	return kernels.GemvQ5_0F32(wQ5_0, x, y, M, K, streamPtr(s))
+func (k *CUDAKernels) GemvQ5_0F32(wQ5_0, x, y unsafe.Pointer, M, K, qhOffset, qsOffset int, s Stream) error {
+	return kernels.GemvQ5_0F32(wQ5_0, x, y, M, K, qhOffset, qsOffset, streamPtr(s))
 }
 
 func (k *CUDAKernels) DequantQ4KF32(src, dst unsafe.Pointer, rows, K int, s Stream) error {

internal/gpuapi/fpga_kernels.go

@@ -115,7 +115,7 @@ func (k *FPGAKernels) GemvQ6KF32(_, _, _ unsafe.Pointer, _, _ int, _ Stream) err
 	return fmt.Errorf("GemvQ6KF32: not implemented for FPGA")
 }
 
-func (k *FPGAKernels) GemvQ5_0F32(_, _, _ unsafe.Pointer, _, _ int, _ Stream) error {
+func (k *FPGAKernels) GemvQ5_0F32(_, _, _ unsafe.Pointer, _, _, _, _ int, _ Stream) error {
 	return fmt.Errorf("GemvQ5_0F32: not implemented for FPGA")
 }
 

internal/gpuapi/gpuapi_test.go

@@ -150,7 +150,7 @@ func (stubKernelRunner) GemvQ5KF32(_, _, _ unsafe.Pointer, _, _ int, _ gpuapi.St
 func (stubKernelRunner) GemvQ6KF32(_, _, _ unsafe.Pointer, _, _ int, _ gpuapi.Stream) error {
 	return nil
 }
-func (stubKernelRunner) GemvQ5_0F32(_, _, _ unsafe.Pointer, _, _ int, _ gpuapi.Stream) error {
+func (stubKernelRunner) GemvQ5_0F32(_, _, _ unsafe.Pointer, _, _, _, _ int, _ gpuapi.Stream) error {
 	return nil
 }
 func (stubKernelRunner) DequantQ4KF32(_, _ unsafe.Pointer, _, _ int, _ gpuapi.Stream) error {

internal/gpuapi/kernels.go

@@ -63,9 +63,10 @@ type KernelRunner interface {
 	GemvQ6KF32(wQ6K, x, y unsafe.Pointer, M, K int, stream Stream) error
 
 	// GemvQ5_0F32 performs Q5_0 fused dequant-GEMV: y = dequant(W_q5_0) * x.
-	// W_q5_0 is raw Q5_0 blocks for matrix [M, K]. x is [K] float32.
+	// W_q5_0 is the separated GPU layout (scales | qh | qs). x is [K] float32.
 	// y is [M] float32. K must be a multiple of 32. Batch=1 only.
-	GemvQ5_0F32(wQ5_0, x, y unsafe.Pointer, M, K int, stream Stream) error
+	// qhOffset and qsOffset are byte offsets to the qh and qs regions.
+	GemvQ5_0F32(wQ5_0, x, y unsafe.Pointer, M, K, qhOffset, qsOffset int, stream Stream) error
 
 	// DequantQ4KF32 dequantizes Q4_K super-blocks to FP32 in global memory.
 	// src is raw Q4_K super-blocks for matrix [rows, K]. dst is [rows, K] float32.

internal/gpuapi/metal_kernels.go

@@ -491,7 +491,7 @@ func (k *MetalKernels) GemvQ6KF32(_, _, _ unsafe.Pointer, _, _ int, _ Stream) er
 	return fmt.Errorf("GemvQ6KF32: not yet implemented for Metal")
 }
 
-func (k *MetalKernels) GemvQ5_0F32(_, _, _ unsafe.Pointer, _, _ int, _ Stream) error {
+func (k *MetalKernels) GemvQ5_0F32(_, _, _ unsafe.Pointer, _, _, _, _ int, _ Stream) error {
 	return fmt.Errorf("GemvQ5_0F32: not yet implemented for Metal")
 }