From 15e755f30d3deb75fea2c9cc8cd129a8d7cf634d Mon Sep 17 00:00:00 2001 From: David Friehs Date: Fri, 31 Jul 2026 10:15:44 +0200 Subject: [PATCH] cuda: extract Q2_0 elements via __byte_perm (#25603) --- ggml/src/ggml-cuda/mmq-load-tiles.cuh | 48 +++++++++---------------- ggml/src/ggml-cuda/vecdotq.cuh | 52 +++++++++------------------ 2 files changed, 32 insertions(+), 68 deletions(-) diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index 8c9070c6be..8ed704c281 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -130,43 +130,27 @@ template static __device__ __forceinline_ } const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + kbx; - // Each 32-element chunk occupies 8 bytes of qs (32 elements * 2 bits = 64 bits) - const int qs_offset = 8*kqsx; - const int qs0 = bxi->qs[qs_offset + 0] | (bxi->qs[qs_offset + 1] << 8) | - (bxi->qs[qs_offset + 2] << 16) | (bxi->qs[qs_offset + 3] << 24); - const int qs1 = bxi->qs[qs_offset + 4] | (bxi->qs[qs_offset + 5] << 8) | - (bxi->qs[qs_offset + 6] << 16) | (bxi->qs[qs_offset + 7] << 24); - - // Unpack 32 2-bit codes into 8 int32s, each holding 4 signed int8s in {-1,0,1,2}. - int unpacked_bytes[8]; -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int shift = j * 8; - const int codes = (qs0 >> shift) & 0xFF; - const int c0 = ((codes >> 0) & 0x3) - 1; - const int c1 = ((codes >> 2) & 0x3) - 1; - const int c2 = ((codes >> 4) & 0x3) - 1; - const int c3 = ((codes >> 6) & 0x3) - 1; - unpacked_bytes[j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24); - } -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int shift = j * 8; - const int codes = (qs1 >> shift) & 0xFF; - const int c0 = ((codes >> 0) & 0x3) - 1; - const int c1 = ((codes >> 2) & 0x3) - 1; - const int c2 = ((codes >> 4) & 0x3) - 1; - const int c3 = ((codes >> 6) & 0x3) - 1; - unpacked_bytes[4 + j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24); - } + const int16_t * qxi = (const int16_t *) bxi->qs + kqsx * 4; const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0; + #pragma unroll - for (int j = 0; j < 8; ++j) { + for (int j = 0; j < 4; ++j) { + const int q = qxi[j]; + + // unpack even and odd crumbs into byte values + const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); + const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); + // unshuffle values + const int qx = __byte_perm(qe, qo, 0x5140); + const int qy = __byte_perm(qe, qo, 0x7362); + #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) - x_qs[i*sram_stride + dst_offset + j] = unpacked_bytes[j]; + x_qs[i*sram_stride + dst_offset + j*2+0] = qx; + x_qs[i*sram_stride + dst_offset + j*2+1] = qy; #else - x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j] = unpacked_bytes[j]; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+0] = qx; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+1] = qy; #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) } } diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index 36d89c89e3..0f039c735b 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -734,48 +734,28 @@ static __device__ __forceinline__ float vec_dot_q2_0_q8_1( // Q8_1: 32 elements per block with individual scales // iqs selects which of the 2 chunks of 32 elements to process (0-1) - const float d2 = bq2_0->d; + const float d2 = bq2_0->d; + const int16_t * qs = (const int16_t *) bq2_0->qs + iqs * 4; // Process only the chunk specified by iqs const block_q8_1 * bq8_1_chunk = bq8_1 + iqs; - // Load 64 bits (8 bytes) for this chunk from Q2_0: bytes [8*iqs, 8*iqs+8) - const int offset = iqs * 8; - const int v0 = bq2_0->qs[offset + 0] | (bq2_0->qs[offset + 1] << 8) | - (bq2_0->qs[offset + 2] << 16) | (bq2_0->qs[offset + 3] << 24); - const int v1 = bq2_0->qs[offset + 4] | (bq2_0->qs[offset + 5] << 8) | - (bq2_0->qs[offset + 6] << 16) | (bq2_0->qs[offset + 7] << 24); - - // Unpack 32 2-bit codes into 8 int32s, each holding 4 signed int8 symbols in {-1,0,1,2}. - // Stored code c in {0,1,2,3} -> symbol s = c - 1. - int vi_bytes[8]; -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int shift = j * 8; - const int codes = (v0 >> shift) & 0xFF; - const int c0 = ((codes >> 0) & 0x3) - 1; - const int c1 = ((codes >> 2) & 0x3) - 1; - const int c2 = ((codes >> 4) & 0x3) - 1; - const int c3 = ((codes >> 6) & 0x3) - 1; - vi_bytes[j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24); - } -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int shift = j * 8; - const int codes = (v1 >> shift) & 0xFF; - const int c0 = ((codes >> 0) & 0x3) - 1; - const int c1 = ((codes >> 2) & 0x3) - 1; - const int c2 = ((codes >> 4) & 0x3) - 1; - const int c3 = ((codes >> 6) & 0x3) - 1; - vi_bytes[4 + j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24); - } - - // Compute dot product for this 32-element chunk int sumi = 0; #pragma unroll - for (int j = 0; j < 8; ++j) { - const int u = get_int_b4(bq8_1_chunk->qs, j); - sumi = ggml_cuda_dp4a(vi_bytes[j], u, sumi); + for (int j = 0; j < 4; ++j) { + const int q = qs[j]; + const int u = get_int_b4(bq8_1_chunk->qs, j*2+0); + const int v = get_int_b4(bq8_1_chunk->qs, j*2+1); + + // unpack even and odd crumbs into byte values + const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); + const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); + // unshuffle values + const int qx = __byte_perm(qe, qo, 0x5140); + const int qy = __byte_perm(qe, qo, 0x7362); + + sumi = ggml_cuda_dp4a(u, qx, sumi); + sumi = ggml_cuda_dp4a(v, qy, sumi); } // Apply Q2_0's single scale and this chunk's Q8_1 scale