Skip to content

cuda: extract Q2_0 elements via __byte_perm #104

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
khosravipasha merged 1 commit into from
Jul 31, 2026
Merged

cuda: extract Q2_0 elements via __byte_perm #104

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 16 additions & 31 deletions ggml/src/ggml-cuda/mmq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -425,43 +425,28 @@ template <int mmq_y, bool need_check> static __device__ __forceinline__ void loa
}

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);
// each 32-element chunk occupies 8 bytes of qs (4 int16), regardless of group size
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;

// 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 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);

const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
#pragma unroll
for (int j = 0; j < 8; ++j) {
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + dst_offset + j] = unpacked_bytes[j];
x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + dst_offset + j*2+0] = qx;
x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 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)
}
}
Expand Down
65 changes: 21 additions & 44 deletions ggml/src/ggml-cuda/vecdotq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -746,58 +746,35 @@ static __device__ __forceinline__ float vec_dot_q2_0_q8_1(
// Q8_1: 32 elements per block with individual scales
// iqs selects which of the 4 chunks of 32 elements to process (0-3)

const float d2 = bq2_0->d;
const float d2 = bq2_0->d;
// each 32-element chunk occupies 8 bytes of qs (4 int16), regardless of group size
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 of raw UNSIGNED codes {0,1,2,(3)} --
// no per-element "-1" offset. Symbol s = code - 1, so sum(s*act) =
// sum(code*act) - sum(act); that correction is applied once at the end
// instead (matches the deferred-correction pattern vec_dot_q4_0_q8_1_impl
// uses -- code 3 is unreachable from the reference quantizer, so this
// covers the only codes {0,1,2} that ever actually occur).
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;
const int c1 = (codes >> 2) & 0x3;
const int c2 = (codes >> 4) & 0x3;
const int c3 = (codes >> 6) & 0x3;
vi_bytes[j] = c0 | (c1 << 8) | (c2 << 16) | (c3 << 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;
const int c1 = (codes >> 2) & 0x3;
const int c2 = (codes >> 4) & 0x3;
const int c3 = (codes >> 6) & 0x3;
vi_bytes[4 + j] = c0 | (c1 << 8) | (c2 << 16) | (c3 << 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);

// the permute source 0x020100FF holds bytes {-1,0,1,2}, so a 2-bit code
// indexes straight to its symbol (s = code - 1) with no arithmetic
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);
}

// ds.x = d8 (per-block activation scale), ds.y = sum(act) in real units
// (see quantize_q8_1: y[ib].ds = make_half2(d, sum)).
const float2 ds8f = __half22float2(bq8_1_chunk->ds);
return d2 * (sumi * ds8f.x - ds8f.y);
// symbols are already signed, so no deferred sum(act) correction is needed
const float d8 = __low2float(bq8_1_chunk->ds);
return d2 * d8 * sumi;
}

static __device__ __forceinline__ float vec_dot_q4_0_q8_1(
Expand Down
Loading