Skip to content

vulkan: optimize iq3_xxs mat-vec shader #28427

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

Closed
WakeUpMorty wants to merge 1 commit into from
+18 −12
Closed

vulkan: optimize iq3_xxs mat-vec shader #28427

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
30 changes: 18 additions & 12 deletions ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_iq3_xxs.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,15 @@ layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;

FLOAT_TYPE temp[NUM_COLS][NUM_ROWS];

// invocations per superblock. with many columns, 8 invocations need too many
// registers and spill, so use 16 to halve the per-invocation B working set
const uint TPB = NUM_COLS <= 4 ? 8 : 16;
const uint NL = 32 / TPB; // l steps per invocation

void calc_superblock(const uint a_offset, const uint b_offset, const uint itid, const uint i, const uint num_blocks_per_row, const uint first_row, const uint num_rows) {
const uint y_idx = i * QUANT_K + 16 * itid;
const uint ib32 = itid / 2; // 0..7
const uint ib32 = itid / (TPB / 8);
const uint l0 = (itid % (TPB / 8)) * NL;
const uint y_idx = i * QUANT_K + 32 * ib32;

uint ibi = a_offset + first_row * num_blocks_per_row + i;
[[unroll]] for (uint n = 0; n < num_rows; ++n) {
Expand All @@ -18,13 +24,13 @@ void calc_superblock(const uint a_offset, const uint b_offset, const uint itid,
data_a_packed16[ibi].qs[QUANT_K / 8 + 2 * ib32],
data_a_packed16[ibi].qs[QUANT_K / 8 + 2 * ib32 + 1]));
const float db = d * 0.5 * (0.5 + (signscale >> 28));
[[unroll]] for (uint l = 0; l < 2; ++l) {
const uint qs0 = data_a[ibi].qs[8 * ib32 + 4 * (itid & 1) + 2 * l];
const uint qs1 = data_a[ibi].qs[8 * ib32 + 4 * (itid & 1) + 2 * l + 1];
const uint sign = bitfieldExtract(signscale, 7 * int(2 * (itid & 1) + l), 7);
[[unroll]] for (uint ll = 0; ll < NL; ++ll) {
const uint l = l0 + ll;
const uint qs = data_a_packed16[ibi].qs[4 * ib32 + l]; // one u16 per l: low + high grid index
const uint sign = bitfieldExtract(signscale, 7 * int(l), 7);
const uint sign7 = bitCount(sign);
const vec4 grid0 = vec4(unpack8(iq3xxs_grid[qs0]));
const vec4 grid1 = vec4(unpack8(iq3xxs_grid[qs1]));
const vec4 grid0 = vec4(unpack8(iq3xxs_grid[qs & 0xFF]));
const vec4 grid1 = vec4(unpack8(iq3xxs_grid[qs >> 8]));

[[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
const vec4 b0 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 2*l + 0]);
Expand Down Expand Up @@ -53,11 +59,11 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {

const uint num_blocks_per_row = p.ncols / QUANT_K;

// 16 threads are used to process each block
const uint blocks_per_wg = gl_WorkGroupSize.x/16;
// TPB invocations are used to process each block
const uint blocks_per_wg = gl_WorkGroupSize.x/TPB;
const uint tid = gl_LocalInvocationID.x;
const uint itid = tid % 16; // 0...15
const uint ix = tid / 16;
const uint itid = tid % TPB;
const uint ix = tid / TPB;

[[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
[[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) {
Expand Down