happyz
/
llama.cpp
mirror of https://github.com/ggerganov/llama.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

vulkan: handle large sizes for get_rows (#15686)

This commit is contained in:
Jeff Bolz 2025-08-31 03:13:27 -05:00 committed by GitHub
parent c37052ab4d
commit bbbf5ecccb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 46 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ggml/src/ggml-vulkan/ggml-vulkan.cpp
View File

@ -7849,6 +7849,8 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
break; break;
case GGML_OP_GET_ROWS: case GGML_OP_GET_ROWS:
elements = { (uint32_t)ne00, (uint32_t)ne10, (uint32_t)(ne11 * ne12) }; elements = { (uint32_t)ne00, (uint32_t)ne10, (uint32_t)(ne11 * ne12) };
elements[1] = std::min(elements[1], ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
elements[2] = std::min(elements[2], ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
break; break;
case GGML_OP_ARGSORT: case GGML_OP_ARGSORT:
elements = { (uint32_t)ne00, (uint32_t)ggml_nrows(src0), 1 }; elements = { (uint32_t)ne00, (uint32_t)ggml_nrows(src0), 1 };

29
ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp
View File

@ -7,27 +7,36 @@ layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
void main() { void main() {
const uint i00 = gl_GlobalInvocationID.x; const uint i00 = gl_GlobalInvocationID.x;
const uint i10 = gl_GlobalInvocationID.y;
const uint i11 = (gl_GlobalInvocationID.z)/p.ne12;
const uint i12 = (gl_GlobalInvocationID.z)%p.ne12;
if (i00 >= p.ne00) { if (i00 >= p.ne00) {
return; return;
} }
const uint i01 = data_b[get_boffset() + i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; uint gid_z = gl_GlobalInvocationID.z;
while (gid_z < p.ne11 * p.ne12) {
uint gid_y = gl_GlobalInvocationID.y;
while (gid_y < p.ne10) {
const uint i10 = gid_y;
const uint i11 = gid_z / p.ne12;
const uint i12 = gid_z % p.ne12;
const uint a_offset = get_aoffset() + i01*p.nb01 + i11*p.nb02 + i12*p.nb03; const uint i01 = data_b[get_boffset() + i10*p.nb10 + i11*p.nb11 + i12*p.nb12];
const uint d_offset = get_doffset() + i10*p.nb21 + i11*p.nb22 + i12*p.nb23;
const uint a_offset = get_aoffset() + i01*p.nb01 + i11*p.nb02 + i12*p.nb03;
const uint d_offset = get_doffset() + i10*p.nb21 + i11*p.nb22 + i12*p.nb23;
#if defined(DATA_A_BF16) #if defined(DATA_A_BF16)
FLOAT_TYPE v = FLOAT_TYPE(bf16_to_fp32(data_a[a_offset + i00])); FLOAT_TYPE v = FLOAT_TYPE(bf16_to_fp32(data_a[a_offset + i00]));
#else #else
FLOAT_TYPE v = FLOAT_TYPE(data_a[a_offset + i00]); FLOAT_TYPE v = FLOAT_TYPE(data_a[a_offset + i00]);
#endif #endif
#ifndef OPTIMIZATION_ERROR_WORKAROUND #ifndef OPTIMIZATION_ERROR_WORKAROUND
data_d[d_offset + i00] = D_TYPE(v); data_d[d_offset + i00] = D_TYPE(v);
#else #else
data_d[d_offset + i00] = D_TYPE(v); data_d[d_offset + i00] = D_TYPE(v);
#endif #endif
gid_y += gl_WorkGroupSize.y * gl_NumWorkGroups.y;
}
gid_z += gl_WorkGroupSize.z * gl_NumWorkGroups.z;
}
} }

40
ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp
View File

@ -10,9 +10,6 @@ layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
void main() { void main() {
const uint i00 = (gl_GlobalInvocationID.x)*2; const uint i00 = (gl_GlobalInvocationID.x)*2;
const uint i10 = gl_GlobalInvocationID.y;
const uint i11 = (gl_GlobalInvocationID.z)/p.ne12;
const uint i12 = (gl_GlobalInvocationID.z)%p.ne12;
#ifdef NEEDS_INIT_IQ_SHMEM #ifdef NEEDS_INIT_IQ_SHMEM
init_iq_shmem(gl_WorkGroupSize); init_iq_shmem(gl_WorkGroupSize);
@ -22,20 +19,33 @@ void main() {
return; return;
} }
const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; uint gid_z = gl_GlobalInvocationID.z;
while (gid_z < p.ne11 * p.ne12) {
uint gid_y = gl_GlobalInvocationID.y;
while (gid_y < p.ne10) {
const uint i10 = gid_y;
const uint i11 = gid_z / p.ne12;
const uint i12 = gid_z % p.ne12;
const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03; const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12];
const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23;
const uint ib = a_offset + i00/QUANT_K; // block index const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03;
const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23;
const uint iybs = i00 - i00%QUANT_K; // dst block start index
const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
vec2 v = dequantize(ib, iqs, 0); const uint ib = a_offset + i00/QUANT_K; // block index
const vec2 dm = get_dm(ib, 0); const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index
v = v * dm.x + dm.y; const uint iybs = i00 - i00%QUANT_K; // dst block start index
const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
data_d[d_offset + iybs + iqs ] = D_TYPE(v.x); vec2 v = dequantize(ib, iqs, 0);
data_d[d_offset + iybs + iqs + y_offset] = D_TYPE(v.y); const vec2 dm = get_dm(ib, 0);
v = v * dm.x + dm.y;
data_d[d_offset + iybs + iqs ] = D_TYPE(v.x);
data_d[d_offset + iybs + iqs + y_offset] = D_TYPE(v.y);
gid_y += gl_WorkGroupSize.y * gl_NumWorkGroups.y;
}
gid_z += gl_WorkGroupSize.z * gl_NumWorkGroups.z;
}
} }