From f1eb80ef5e989b04108984b48b8b19e549a626cb Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Wed, 25 Mar 2026 15:21:22 -0700 Subject: [PATCH 01/16] Work towards removing bitcast --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 30 ++++- .../ggml-webgpu/wgsl-shaders/flash_attn.wgsl | 80 ++++++++++---- .../wgsl-shaders/mul_mat_decls.tmpl | 103 ++++++++++-------- .../ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl | 97 ++++++++++------- 4 files changed, 202 insertions(+), 108 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index 59861ac16cc3..c24357be6899 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -1125,8 +1125,18 @@ class ggml_webgpu_shader_lib { defines.push_back("BYTE_HELPERS"); defines.push_back("MUL_ACC_" + type_upper); - // For fast path we always dequantize from f16 inside the shader - defines.push_back("SRC0_INNER_TYPE=f16"); + if (context.src0->type == GGML_TYPE_Q4_0 || + context.src0->type == GGML_TYPE_Q4_1 || + context.src0->type == GGML_TYPE_Q5_0 || + context.src0->type == GGML_TYPE_Q5_1 || + context.src0->type == GGML_TYPE_Q8_0 || + context.src0->type == GGML_TYPE_Q8_1) { + defines.push_back("RAW_SRC0_BYTES"); + defines.push_back("SRC0_INNER_TYPE=u32"); + } else { + // For fast path we always dequantize from f16 inside the shader + defines.push_back("SRC0_INNER_TYPE=f16"); + } break; } } @@ -1240,8 +1250,18 @@ class ggml_webgpu_shader_lib { defines.push_back("INIT_SRC0_SHMEM_" + type_upper); defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - // Use f16 inside the shader for quantized types - defines.push_back("SRC0_INNER_TYPE=f16"); + if (context.src0->type == GGML_TYPE_Q4_0 || + context.src0->type == GGML_TYPE_Q4_1 || + context.src0->type == GGML_TYPE_Q5_0 || + context.src0->type == GGML_TYPE_Q5_1 || + context.src0->type == GGML_TYPE_Q8_0 || + context.src0->type == GGML_TYPE_Q8_1) { + defines.push_back("RAW_SRC0_BYTES"); + defines.push_back("SRC0_INNER_TYPE=u32"); + } else { + // Use f16 inside the shader for quantized types + defines.push_back("SRC0_INNER_TYPE=f16"); + } variant += std::string("_") + src0_name; break; @@ -1614,9 +1634,11 @@ class ggml_webgpu_shader_lib { break; case GGML_TYPE_Q4_0: defines.push_back("KV_Q4_0"); + defines.push_back("KV_RAW_U32"); break; case GGML_TYPE_Q8_0: defines.push_back("KV_Q8_0"); + defines.push_back("KV_RAW_U32"); break; default: GGML_ABORT("Unsupported KV type for flash attention shader"); diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl index b6822161464c..a721c0325a37 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl @@ -6,6 +6,8 @@ enable chromium_experimental_subgroup_matrix; #ifdef KV_F32 #define KV_TYPE f32 +#elif defined(KV_RAW_U32) +#define KV_TYPE u32 #else #define KV_TYPE f16 #endif @@ -37,11 +39,13 @@ enable chromium_experimental_subgroup_matrix; #define NQ 16 // Q4_0 has 32 elements, 1 f16 for scale, 8 f16 for 4-bit weights #define F16_PER_BLOCK 9 +#define BLOCK_SIZE_BYTES 18u #define WEIGHTS_PER_F16 4 #elif defined(KV_Q8_0) #define NQ 8 // Q8_0 has 32 elements, 1 f16 for scale, 16 f16 for 8-bit weights #define F16_PER_BLOCK 17 +#define BLOCK_SIZE_BYTES 34u #define WEIGHTS_PER_F16 2 #endif #define F16_PER_THREAD (NQ / WEIGHTS_PER_F16) @@ -55,6 +59,46 @@ fn get_byte_i32(value: u32, index: u32) -> i32 { return bitcast(((value >> (index * 8)) & 0xFF) << 24) >> 24; } +#ifdef KV_RAW_U32 +fn load_k_u16_at(byte_offset: u32) -> u32 { + let word = K[byte_offset / 4u]; + let shift = (byte_offset & 2u) * 8u; + return (word >> shift) & 0xFFFFu; +} + +fn load_k_u32_at(byte_offset: u32) -> u32 { + let word_idx = byte_offset / 4u; + let shift = (byte_offset & 3u) * 8u; + let lo = K[word_idx]; + if (shift == 0u) { + return lo; + } + let hi = K[word_idx + 1u]; + return (lo >> shift) | (hi << (32u - shift)); +} + +fn load_v_u16_at(byte_offset: u32) -> u32 { + let word = V[byte_offset / 4u]; + let shift = (byte_offset & 2u) * 8u; + return (word >> shift) & 0xFFFFu; +} + +fn load_v_u32_at(byte_offset: u32) -> u32 { + let word_idx = byte_offset / 4u; + let shift = (byte_offset & 3u) * 8u; + let lo = V[word_idx]; + if (shift == 0u) { + return lo; + } + let hi = V[word_idx + 1u]; + return (lo >> shift) | (hi << (32u - shift)); +} + +fn f16_from_u16(bits: u32) -> f16 { + return bitcast>(bits).x; +} +#endif + struct Params { offset_q: u32, offset_k: u32, @@ -254,12 +298,11 @@ fn main(@builtin(workgroup_id) wg_id: vec3, if (global_k_row < params.seq_len_kv) { let global_block_idx = k_head_offset + global_k_row * params.stride_k1 + block_k; - let base_idx = global_block_idx * F16_PER_BLOCK; - let d = K[base_idx]; // scale + let block_byte_base = global_block_idx * BLOCK_SIZE_BYTES; + let d = f16_from_u16(load_k_u16_at(block_byte_base)); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = K[base_idx + 1u + block_offset + j]; - let q_1 = K[base_idx + 1u + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_k_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte(q_packed, k); let q_hi = (f16((q_byte >> 4) & 0xF) - 8.0) * d; @@ -282,12 +325,11 @@ fn main(@builtin(workgroup_id) wg_id: vec3, if (global_k_row < params.seq_len_kv) { let global_block_idx = k_head_offset + global_k_row * params.stride_k1 + block_k; - let base_idx = global_block_idx * F16_PER_BLOCK; - let d = K[base_idx]; // scale + let block_byte_base = global_block_idx * BLOCK_SIZE_BYTES; + let d = f16_from_u16(load_k_u16_at(block_byte_base)); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = K[base_idx + 1u + block_offset + j]; - let q_1 = K[base_idx + 1u + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_k_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte_i32(q_packed, k); let q_val = f16(q_byte) * d; @@ -459,12 +501,11 @@ fn main(@builtin(workgroup_id) wg_id: vec3, if (global_v_row < params.seq_len_kv) { let global_block_idx = v_head_offset + global_v_row * params.stride_v1 + block_k; - let base_idx = global_block_idx * F16_PER_BLOCK; - let d = V[base_idx]; // scale + let block_byte_base = global_block_idx * BLOCK_SIZE_BYTES; + let d = f16_from_u16(load_v_u16_at(block_byte_base)); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = V[base_idx + 1u + block_offset + j]; - let q_1 = V[base_idx + 1u + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_v_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte(q_packed, k); let q_hi = (f16((q_byte >> 4) & 0xF) - 8.0) * d; @@ -487,12 +528,11 @@ fn main(@builtin(workgroup_id) wg_id: vec3, if (global_v_row < params.seq_len_kv) { let global_block_idx = v_head_offset + global_v_row * params.stride_v1 + block_k; - let base_idx = global_block_idx * F16_PER_BLOCK; - let d = V[base_idx]; // scale + let block_byte_base = global_block_idx * BLOCK_SIZE_BYTES; + let d = f16_from_u16(load_v_u16_at(block_byte_base)); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = V[base_idx + 1u + block_offset + j]; - let q_1 = V[base_idx + 1u + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_v_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte_i32(q_packed, k); let q_val = f16(q_byte) * d; diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl index de60ebbcf2b5..233333b86789 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl @@ -25,6 +25,29 @@ fn store_shmem(val: f16, idx: u32) { } #endif // SCALAR +#ifdef RAW_SRC0_BYTES +fn load_src0_u16_at(byte_offset: u32) -> u32 { + let word = src0[byte_offset / 4u]; + let shift = (byte_offset & 2u) * 8u; + return (word >> shift) & 0xFFFFu; +} + +fn load_src0_u32_at(byte_offset: u32) -> u32 { + let word_idx = byte_offset / 4u; + let shift = (byte_offset & 3u) * 8u; + let lo = src0[word_idx]; + if (shift == 0u) { + return lo; + } + let hi = src0[word_idx + 1u]; + return (lo >> shift) | (hi << (32u - shift)); +} + +fn load_src0_f16_at(byte_offset: u32) -> f16 { + return bitcast>(load_src0_u16_at(byte_offset)).x; +} +#endif // RAW_SRC0_BYTES + #ifdef INIT_SRC0_SHMEM_FLOAT fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { for (var elem_idx = thread_id * VEC_SIZE; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE * VEC_SIZE) { @@ -61,6 +84,7 @@ fn init_shmem_src1(thread_id: u32, batch_offset: u32, offset_n: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q4_0 const BLOCK_SIZE = 32u; +const BLOCK_SIZE_BYTES = 18u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; @@ -81,14 +105,12 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 if (global_m < params.m && global_k < params.k / BLOCK_SIZE) { let src0_idx = batch_offset + global_m * params.stride_01 + global_k; - let scale_idx = src0_idx * F16_PER_BLOCK; - let d = src0[scale_idx]; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; + let d = load_src0_f16_at(block_byte_base); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = src0[scale_idx + 1u + block_offset + j]; - let q_1 = src0[scale_idx + 1u + block_offset + j + 1]; - - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte(q_packed, k); let q_hi = (f16((q_byte >> 4) & 0xF) - 8.0) * d; @@ -104,6 +126,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q4_1 const BLOCK_SIZE = 32u; +const BLOCK_SIZE_BYTES = 20u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; @@ -124,15 +147,13 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 if (global_m < params.m && global_k < params.k / BLOCK_SIZE) { let src0_idx = batch_offset + global_m * params.stride_01 + global_k; - let scale_idx = src0_idx * F16_PER_BLOCK; - let d = src0[scale_idx]; - let m = src0[scale_idx + 1u]; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; + let d = load_src0_f16_at(block_byte_base); + let m = load_src0_f16_at(block_byte_base + 2u); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = src0[scale_idx + 2u + block_offset + j]; - let q_1 = src0[scale_idx + 2u + block_offset + j + 1]; - - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 4u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte(q_packed, k); let q_lo = f16(q_byte & 0xF) * d + m; @@ -149,6 +170,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q5_0 // 32 weights per block, each at 4 bits each = 32 * 4 = 128 bits / 16 = 8 f16s per block const BLOCK_SIZE = 32u; +const BLOCK_SIZE_BYTES = 22u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. // tile_k is defined as 32u, so blocks_k ends up being 1 always override BLOCKS_K = TILE_K / BLOCK_SIZE; @@ -171,18 +193,14 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 if (global_m < params.m && global_k < params.k / BLOCK_SIZE) { let src0_idx = batch_offset + global_m * params.stride_01 + global_k; - let scale_idx = src0_idx * F16_PER_BLOCK; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; - let d = src0[scale_idx]; - let qh0 = src0[scale_idx + 1u]; - let qh1 = src0[scale_idx + 2u]; - let qh_packed = bitcast(vec2(qh0, qh1)); + let d = load_src0_f16_at(block_byte_base); + let qh_packed = load_src0_u32_at(block_byte_base + 2u); for (var j = 0u; j < 2; j++) { - let q_0 = src0[scale_idx + 3u + block_offset + (j*2)]; - let q_1 = src0[scale_idx + 3u + block_offset + (j*2) + 1u]; - - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 6u + 2u * (block_offset + j * 2u); + let q_packed = load_src0_u32_at(q_byte_offset); let j_adjusted = j + (block_offset / 2u); @@ -207,6 +225,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q5_1 // 32 weights per block, each at 4 bits each = 32 * 4 = 128 bits / 16 = 8 f16s per block const BLOCK_SIZE = 32u; +const BLOCK_SIZE_BYTES = 24u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. // tile_k is defined as 32u, so blocks_k ends up being 1 always override BLOCKS_K = TILE_K / BLOCK_SIZE; @@ -229,20 +248,16 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 if (global_m < params.m && global_k < params.k / BLOCK_SIZE) { let src0_idx = batch_offset + global_m * params.stride_01 + global_k; - let scale_idx = src0_idx * F16_PER_BLOCK; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; - let d = src0[scale_idx]; - let m = src0[scale_idx + 1u]; - let qh0 = src0[scale_idx + 2u]; - let qh1 = src0[scale_idx + 3u]; - let qh_packed = bitcast(vec2(qh0, qh1)); + let d = load_src0_f16_at(block_byte_base); + let m = load_src0_f16_at(block_byte_base + 2u); + let qh_packed = load_src0_u32_at(block_byte_base + 4u); for (var j = 0u; j < 2; j++) { - let q_0 = src0[scale_idx + 4u + block_offset + (j*2)]; - let q_1 = src0[scale_idx + 4u + block_offset + (j*2) + 1u]; - - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 8u + 2u * (block_offset + j * 2u); + let q_packed = load_src0_u32_at(q_byte_offset); let j_adjusted = j + (block_offset / 2u); @@ -266,6 +281,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q8_0 const BLOCK_SIZE = 32u; +const BLOCK_SIZE_BYTES = 34u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; @@ -286,14 +302,12 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 if (global_m < params.m && global_k < params.k / BLOCK_SIZE) { let src0_idx = batch_offset + global_m * params.stride_01 + global_k; - let scale_idx = src0_idx * F16_PER_BLOCK; - let d = src0[scale_idx]; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; + let d = load_src0_f16_at(block_byte_base); for (var j = 0u; j < F16_PER_THREAD; j+=2) { - let q_0 = src0[scale_idx + 1u + block_offset + j]; - let q_1 = src0[scale_idx + 1u + block_offset + j + 1]; - - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte_i32(q_packed, k); @@ -308,6 +322,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q8_1 const BLOCK_SIZE = 32u; +const BLOCK_SIZE_BYTES = 36u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; @@ -328,15 +343,13 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 if (global_m < params.m && global_k < params.k / BLOCK_SIZE) { let src0_idx = batch_offset + global_m * params.stride_01 + global_k; - let scale_idx = src0_idx * F16_PER_BLOCK; - let d = src0[scale_idx]; - let m = src0[scale_idx + 1u]; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; + let d = load_src0_f16_at(block_byte_base); + let m = load_src0_f16_at(block_byte_base + 2u); for (var j = 0u; j < F16_PER_THREAD; j+=2) { - let q_0 = src0[scale_idx + 2u + block_offset + j]; - let q_1 = src0[scale_idx + 2u + block_offset + j + 1]; - - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 4u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k = 0u; k < 4u; k++) { let q_byte = get_byte_i32(q_packed, k); diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl index 94f4bae11f4a..697126049062 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl @@ -37,6 +37,29 @@ fn store_val(group_base: u32) -> f32 { } #endif +#ifdef RAW_SRC0_BYTES +fn load_src0_u16_at(byte_offset: u32) -> u32 { + let word = src0[byte_offset / 4u]; + let shift = (byte_offset & 2u) * 8u; + return (word >> shift) & 0xFFFFu; +} + +fn load_src0_u32_at(byte_offset: u32) -> u32 { + let word_idx = byte_offset / 4u; + let shift = (byte_offset & 3u) * 8u; + let lo = src0[word_idx]; + if (shift == 0u) { + return lo; + } + let hi = src0[word_idx + 1u]; + return (lo >> shift) | (hi << (32u - shift)); +} + +fn load_src0_f16_at(byte_offset: u32) -> f16 { + return bitcast>(load_src0_u16_at(byte_offset)).x; +} +#endif + #ifdef MUL_ACC_FLOAT fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { var local_sum = 0.0; @@ -52,6 +75,7 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { #ifdef MUL_ACC_Q4_0 const BLOCK_SIZE = 32; +const BLOCK_SIZE_BYTES = 18u; const NQ = 16u; // number of weights per thread const F16_PER_BLOCK = 9u; // 1 scale + 8x4 packed weights const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 @@ -62,14 +86,13 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { for (var i = tig * NQ; i < tile_size; i += THREADS_PER_OUTPUT * NQ) { let blck_idx = i / BLOCK_SIZE; let block_offset = (i % BLOCK_SIZE) / WEIGHTS_PER_F16; - let scale_idx = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * F16_PER_BLOCK; + let block_byte_base = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * BLOCK_SIZE_BYTES; // each f16 contains offsets [block_offset, block_offset + 1] and [block_offset + 16, block_offset + 17] let shmem_idx = blck_idx * BLOCK_SIZE + block_offset * 2u; - let d = f32(src0[scale_idx]); + let d = f32(load_src0_f16_at(block_byte_base)); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = src0[scale_idx + 1 + block_offset + j]; - let q_1 = src0[scale_idx + 1 + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k: u32 = 0; k < 4; k++) { let q_byte = get_byte(q_packed, k); let q_hi = (f32((q_byte >> 4) & 0xF) - 8.0) * d; @@ -86,6 +109,7 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { #ifdef MUL_ACC_Q4_1 const BLOCK_SIZE = 32; +const BLOCK_SIZE_BYTES = 20u; const NQ = 16u; // number of weights per thread const F16_PER_BLOCK = 10u; const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 @@ -96,15 +120,14 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { for (var i = tig * NQ; i < tile_size; i += THREADS_PER_OUTPUT * NQ) { let blck_idx = i / BLOCK_SIZE; let block_offset = (i % BLOCK_SIZE) / WEIGHTS_PER_F16; - let scale_idx = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * F16_PER_BLOCK; + let block_byte_base = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * BLOCK_SIZE_BYTES; // each f16 contains offsets [block_offset, block_offset + 1] and [block_offset + 16, block_offset + 17] let shmem_idx = blck_idx * BLOCK_SIZE + block_offset * 2u; - let d = f32(src0[scale_idx]); - let m = f32(src0[scale_idx + 1u]); + let d = f32(load_src0_f16_at(block_byte_base)); + let m = f32(load_src0_f16_at(block_byte_base + 2u)); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = src0[scale_idx + 2u + block_offset + j]; - let q_1 = src0[scale_idx + 2u + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 4u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k: u32 = 0; k < 4; k++) { let q_byte = get_byte(q_packed, k); let q_hi = f32((q_byte >> 4) & 0xF) * d + m; @@ -121,6 +144,7 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { #ifdef MUL_ACC_Q5_0 const BLOCK_SIZE = 32; +const BLOCK_SIZE_BYTES = 22u; const NQ = 16u; // number of weights per thread const F16_PER_BLOCK = 11u; const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 @@ -131,18 +155,15 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { for (var i = tig * NQ; i < tile_size; i += THREADS_PER_OUTPUT * NQ) { let blck_idx = i / BLOCK_SIZE; let block_offset = (i % BLOCK_SIZE) / WEIGHTS_PER_F16; - let scale_idx = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * F16_PER_BLOCK; + let block_byte_base = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * BLOCK_SIZE_BYTES; // each f16 contains offsets [block_offset, block_offset + 1] and [block_offset + 16, block_offset + 17] let shmem_idx = blck_idx * BLOCK_SIZE + block_offset * 2u; - let d = f32(src0[scale_idx]); - let qh0 = src0[scale_idx + 1u]; - let qh1 = src0[scale_idx + 2u]; - let qh_packed = bitcast(vec2(qh0, qh1)); + let d = f32(load_src0_f16_at(block_byte_base)); + let qh_packed = load_src0_u32_at(block_byte_base + 2u); for (var j = 0u; j < 2; j++) { - let q_0 = src0[scale_idx + 3u + block_offset + (j*2)]; - let q_1 = src0[scale_idx + 3u + block_offset + (j*2) + 1u]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 6u + 2u * (block_offset + j * 2u); + let q_packed = load_src0_u32_at(q_byte_offset); let j_adjusted = j + (block_offset / 2u); @@ -168,6 +189,7 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { #ifdef MUL_ACC_Q5_1 const BLOCK_SIZE = 32; +const BLOCK_SIZE_BYTES = 24u; const NQ = 16u; // number of weights per thread const F16_PER_BLOCK = 12u; const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 @@ -178,19 +200,16 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { for (var i = tig * NQ; i < tile_size; i += THREADS_PER_OUTPUT * NQ) { let blck_idx = i / BLOCK_SIZE; let block_offset = (i % BLOCK_SIZE) / WEIGHTS_PER_F16; - let scale_idx = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * F16_PER_BLOCK; + let block_byte_base = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * BLOCK_SIZE_BYTES; // each f16 contains offsets [block_offset, block_offset + 1] and [block_offset + 16, block_offset + 17] let shmem_idx = blck_idx * BLOCK_SIZE + block_offset * 2u; - let d = f32(src0[scale_idx]); - let m = src0[scale_idx + 1u]; - let qh0 = src0[scale_idx + 2u]; - let qh1 = src0[scale_idx + 3u]; - let qh_packed = bitcast(vec2(qh0, qh1)); + let d = f32(load_src0_f16_at(block_byte_base)); + let m = load_src0_f16_at(block_byte_base + 2u); + let qh_packed = load_src0_u32_at(block_byte_base + 4u); for (var j = 0u; j < 2; j++) { - let q_0 = src0[scale_idx + 4u + block_offset + (j*2)]; - let q_1 = src0[scale_idx + 4u + block_offset + (j*2) + 1u]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 8u + 2u * (block_offset + j * 2u); + let q_packed = load_src0_u32_at(q_byte_offset); let j_adjusted = j + (block_offset / 2u); @@ -216,6 +235,7 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { #ifdef MUL_ACC_Q8_0 const BLOCK_SIZE = 32; +const BLOCK_SIZE_BYTES = 34u; const NQ = 16u; // number of weights per thread const F16_PER_BLOCK = 17u; const WEIGHTS_PER_F16 = 2u; @@ -226,15 +246,14 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { for (var i = tig * NQ; i < tile_size; i += THREADS_PER_OUTPUT * NQ) { let blck_idx = i / BLOCK_SIZE; let block_offset = (i % BLOCK_SIZE) / WEIGHTS_PER_F16; - let scale_idx = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * F16_PER_BLOCK; + let block_byte_base = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * BLOCK_SIZE_BYTES; // each f16 contains offsets [block_offset, block_offset + 1] and [block_offset + 16, block_offset + 17] let shmem_idx = blck_idx * BLOCK_SIZE + block_offset * 2u; - let d = f32(src0[scale_idx]); + let d = f32(load_src0_f16_at(block_byte_base)); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = src0[scale_idx + 1 + block_offset + j]; - let q_1 = src0[scale_idx + 1 + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 2u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k: u32 = 0; k < 4; k++) { let q_byte = get_byte_i32(q_packed, k); let q_val = f32(q_byte) * d; @@ -250,6 +269,7 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { #ifdef MUL_ACC_Q8_1 const BLOCK_SIZE = 32; +const BLOCK_SIZE_BYTES = 36u; const NQ = 16u; // number of weights per thread const F16_PER_BLOCK = 18u; const WEIGHTS_PER_F16 = 2u; @@ -260,16 +280,15 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { for (var i = tig * NQ; i < tile_size; i += THREADS_PER_OUTPUT * NQ) { let blck_idx = i / BLOCK_SIZE; let block_offset = (i % BLOCK_SIZE) / WEIGHTS_PER_F16; - let scale_idx = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * F16_PER_BLOCK; + let block_byte_base = (idx_base + k_outer / BLOCK_SIZE + blck_idx) * BLOCK_SIZE_BYTES; // each f16 contains offsets [block_offset, block_offset + 1] and [block_offset + 16, block_offset + 17] let shmem_idx = blck_idx * BLOCK_SIZE + block_offset * 2u; - let d = f32(src0[scale_idx]); - let m = src0[scale_idx + 1u]; + let d = f32(load_src0_f16_at(block_byte_base)); + let m = load_src0_f16_at(block_byte_base + 2u); for (var j = 0u; j < F16_PER_THREAD; j += 2) { - let q_0 = src0[scale_idx + 2u + block_offset + j]; - let q_1 = src0[scale_idx + 2u + block_offset + j + 1]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_byte_offset = block_byte_base + 4u + 2u * (block_offset + j); + let q_packed = load_src0_u32_at(q_byte_offset); for (var k: u32 = 0; k < 4; k++) { let q_byte = get_byte_i32(q_packed, k); let q_val = f32(q_byte) * d + f32(m); From e9af481e087f89f234fae16c372243ca0d7a6c91 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Wed, 25 Mar 2026 15:56:27 -0700 Subject: [PATCH 02/16] Move rest of existing types over --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 12 +- .../wgsl-shaders/mul_mat_decls.tmpl | 110 ++++++------------ .../ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl | 23 ++-- 3 files changed, 54 insertions(+), 91 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index c24357be6899..a3675b8ab058 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -1130,7 +1130,8 @@ class ggml_webgpu_shader_lib { context.src0->type == GGML_TYPE_Q5_0 || context.src0->type == GGML_TYPE_Q5_1 || context.src0->type == GGML_TYPE_Q8_0 || - context.src0->type == GGML_TYPE_Q8_1) { + context.src0->type == GGML_TYPE_Q8_1 || + context.src0->type == GGML_TYPE_Q6_K) { defines.push_back("RAW_SRC0_BYTES"); defines.push_back("SRC0_INNER_TYPE=u32"); } else { @@ -1250,12 +1251,17 @@ class ggml_webgpu_shader_lib { defines.push_back("INIT_SRC0_SHMEM_" + type_upper); defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - if (context.src0->type == GGML_TYPE_Q4_0 || + if (context.src0->type == GGML_TYPE_Q2_K || + context.src0->type == GGML_TYPE_Q3_K || + context.src0->type == GGML_TYPE_Q4_0 || context.src0->type == GGML_TYPE_Q4_1 || context.src0->type == GGML_TYPE_Q5_0 || context.src0->type == GGML_TYPE_Q5_1 || + context.src0->type == GGML_TYPE_Q4_K || + context.src0->type == GGML_TYPE_Q5_K || context.src0->type == GGML_TYPE_Q8_0 || - context.src0->type == GGML_TYPE_Q8_1) { + context.src0->type == GGML_TYPE_Q8_1 || + context.src0->type == GGML_TYPE_Q6_K) { defines.push_back("RAW_SRC0_BYTES"); defines.push_back("SRC0_INNER_TYPE=u32"); } else { diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl index 233333b86789..4e58a892658c 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl @@ -364,7 +364,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q2_K const BLOCK_SIZE = 256u; -const F16_PER_BLOCK = 42u; +const BLOCK_SIZE_BYTES = 84u; fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { // Use standard thread layout instead of lane/row_group @@ -384,10 +384,10 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let k_in_block = global_k % BLOCK_SIZE; let src0_idx = batch_offset + global_m * params.stride_01 + block_k; - let scale_idx = src0_idx * F16_PER_BLOCK; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; - let d = src0[scale_idx + 40u]; - let dmin = src0[scale_idx + 41u]; + let d = load_src0_f16_at(block_byte_base + 80u); + let dmin = load_src0_f16_at(block_byte_base + 82u); // Decode the element at position k_in_block let block_of_32 = k_in_block / 32u; @@ -400,18 +400,14 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let is = k_in_block / 16u; - let sc_0 = src0[scale_idx + 2u * (is / 4u)]; - let sc_1 = src0[scale_idx + 2u * (is / 4u) + 1u]; - let sc_packed = bitcast(vec2(sc_0, sc_1)); + let sc_packed = load_src0_u32_at(block_byte_base + 4u * (is / 4u)); let sc = get_byte(sc_packed, is % 4u); let dl = d * f16(sc & 0xFu); let ml = dmin * f16(sc >> 4u); let q_idx = q_b_idx + k + l; - let q_0 = src0[scale_idx + 8u + 2u * (q_idx / 4u)]; - let q_1 = src0[scale_idx + 8u + 2u * (q_idx / 4u) + 1u]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_packed = load_src0_u32_at(block_byte_base + 16u + 4u * (q_idx / 4u)); let q_byte = get_byte(q_packed, q_idx % 4u); let qs_val = (q_byte >> shift) & 3u; @@ -423,7 +419,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q3_K const BLOCK_SIZE = 256u; -const F16_PER_BLOCK = 55u; +const BLOCK_SIZE_BYTES = 110u; fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { for (var elem_idx = thread_id; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE) { @@ -442,9 +438,9 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let k_in_block = global_k % BLOCK_SIZE; let src0_idx = batch_offset + global_m * params.stride_01 + block_k; - let scale_idx = src0_idx * F16_PER_BLOCK; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; - let d = src0[scale_idx + 54u]; + let d = load_src0_f16_at(block_byte_base + 108u); // Load and unpack scales let kmask1: u32 = 0x03030303u; @@ -452,9 +448,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 var scale_vals: array; for (var i: u32 = 0u; i < 4u; i++) { - let scale_0 = src0[scale_idx + 48u + (2u*i)]; - let scale_1 = src0[scale_idx + 48u + (2u*i) + 1u]; - scale_vals[i] = bitcast(vec2(scale_0, scale_1)); + scale_vals[i] = load_src0_u32_at(block_byte_base + 96u + 4u * i); } var tmp: u32 = scale_vals[2]; @@ -466,16 +460,12 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 // Load hmask and qs arrays var hmask_vals: array; for (var i: u32 = 0u; i < 8u; i++) { - let hmask_0 = src0[scale_idx + (2u*i)]; - let hmask_1 = src0[scale_idx + (2u*i) + 1u]; - hmask_vals[i] = bitcast(vec2(hmask_0, hmask_1)); + hmask_vals[i] = load_src0_u32_at(block_byte_base + 4u * i); } var qs_vals: array; for (var i: u32 = 0u; i < 16u; i++) { - let qs_0 = src0[scale_idx + 16u + (2u*i)]; - let qs_1 = src0[scale_idx + 16u + (2u*i) + 1u]; - qs_vals[i] = bitcast(vec2(qs_0, qs_1)); + qs_vals[i] = load_src0_u32_at(block_byte_base + 32u + 4u * i); } let half = k_in_block / 128u; // 0 or 1 @@ -515,7 +505,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q4_K const BLOCK_SIZE = 256u; -const F16_PER_BLOCK = 72u; +const BLOCK_SIZE_BYTES = 144u; fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { for (var elem_idx = thread_id; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE) { @@ -534,17 +524,15 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let k_in_block = global_k % BLOCK_SIZE; let src0_idx = batch_offset + global_m * params.stride_01 + block_k; - let scale_idx = src0_idx * F16_PER_BLOCK; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; - let d = src0[scale_idx]; - let dmin = src0[scale_idx + 1u]; + let d = load_src0_f16_at(block_byte_base); + let dmin = load_src0_f16_at(block_byte_base + 2u); // Load packed scales var scale_vals: array; for (var i: u32 = 0u; i < 3u; i++) { - let scale_0 = src0[scale_idx + 2u + (2u*i)]; - let scale_1 = src0[scale_idx + 2u + (2u*i) + 1u]; - scale_vals[i] = bitcast(vec2(scale_0, scale_1)); + scale_vals[i] = load_src0_u32_at(block_byte_base + 4u + 4u * i); } // Map k_in_block to loop structure: @@ -580,9 +568,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let ml = dmin * f16(mn); let q_idx = q_b_idx + l; - let q_0 = src0[scale_idx + 8u + 2u * (q_idx / 4u)]; - let q_1 = src0[scale_idx + 8u + 2u * (q_idx / 4u) + 1u]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_packed = load_src0_u32_at(block_byte_base + 16u + 4u * (q_idx / 4u)); let q_byte = get_byte(q_packed, q_idx % 4u); let qs_val = (q_byte >> shift) & 0xFu; @@ -595,7 +581,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q5_K const BLOCK_SIZE = 256u; -const F16_PER_BLOCK = 88u; +const BLOCK_SIZE_BYTES = 176u; fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { for (var elem_idx = thread_id; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE) { @@ -614,17 +600,15 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let k_in_block = global_k % BLOCK_SIZE; let src0_idx = batch_offset + global_m * params.stride_01 + block_k; - let scale_idx = src0_idx * F16_PER_BLOCK; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; - let d = src0[scale_idx]; - let dmin = src0[scale_idx + 1u]; + let d = load_src0_f16_at(block_byte_base); + let dmin = load_src0_f16_at(block_byte_base + 2u); // Load packed scales var scale_vals: array; for (var i: u32 = 0u; i < 3u; i++) { - let scale_0 = src0[scale_idx + 2u + (2u*i)]; - let scale_1 = src0[scale_idx + 2u + (2u*i) + 1u]; - scale_vals[i] = bitcast(vec2(scale_0, scale_1)); + scale_vals[i] = load_src0_u32_at(block_byte_base + 4u + 4u * i); } // The original loop processes elements in groups of 64 @@ -664,15 +648,11 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let ml = dmin * f16(mn); let q_idx = q_b_idx + l; - let q_0 = src0[scale_idx + 24u + 2u * (q_idx / 4u)]; - let q_1 = src0[scale_idx + 24u + 2u * (q_idx / 4u) + 1u]; - let q_packed = bitcast(vec2(q_0, q_1)); + let q_packed = load_src0_u32_at(block_byte_base + 48u + 4u * (q_idx / 4u)); let q_byte = get_byte(q_packed, q_idx % 4u); - let qh_0 = src0[scale_idx + 8u + 2u * (l / 4u)]; - let qh_1 = src0[scale_idx + 8u + 2u * (l / 4u) + 1u]; - let qh_packed = bitcast(vec2(qh_0, qh_1)); + let qh_packed = load_src0_u32_at(block_byte_base + 16u + 4u * (l / 4u)); let qh_byte = get_byte(qh_packed, l % 4u); @@ -688,7 +668,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #ifdef INIT_SRC0_SHMEM_Q6_K const BLOCK_SIZE = 256u; -const F16_PER_BLOCK = 105u; +const BLOCK_SIZE_BYTES = 210u; fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { for (var elem_idx = thread_id; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE) { @@ -707,7 +687,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 let k_in_block = global_k % BLOCK_SIZE; let src0_idx = batch_offset + global_m * params.stride_01 + block_k; - let scale_idx = src0_idx * F16_PER_BLOCK; + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; let half = k_in_block / 128u; let pos_in_half = k_in_block % 128u; @@ -720,30 +700,18 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 // Load only ql13 word needed let ql13_flat = ql_b_idx + l; - let ql13_word = ql13_flat / 4u; - let ql13 = bitcast(vec2( - src0[scale_idx + 2u * ql13_word], - src0[scale_idx + 2u * ql13_word + 1u] - )); - let ql13_b = get_byte(ql13, ql13_flat % 4u); + let ql13 = load_src0_u32_at(block_byte_base + ql13_flat); + let ql13_b = get_byte(ql13, 0u); // Load only ql24 word needed let ql24_flat = ql_b_idx + l + 32u; - let ql24_word = ql24_flat / 4u; - let ql24 = bitcast(vec2( - src0[scale_idx + 2u * ql24_word], - src0[scale_idx + 2u * ql24_word + 1u] - )); - let ql24_b = get_byte(ql24, ql24_flat % 4u); + let ql24 = load_src0_u32_at(block_byte_base + ql24_flat); + let ql24_b = get_byte(ql24, 0u); // Load only qh word needed let qh_flat = qh_b_idx + l; - let qh_word = qh_flat / 4u; - let qh = bitcast(vec2( - src0[scale_idx + 64u + 2u * qh_word], - src0[scale_idx + 64u + 2u * qh_word + 1u] - )); - let qh_b = get_byte(qh, qh_flat % 4u); + let qh = load_src0_u32_at(block_byte_base + 128u + qh_flat); + let qh_b = get_byte(qh, 0u); let q1 = f16((ql13_b & 0xFu) | ((qh_b & 3u) << 4u)) - f16(32.0); let q2 = f16((ql24_b & 0xFu) | (((qh_b >> 2u) & 3u) << 4u)) - f16(32.0); @@ -753,14 +721,10 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 // Load only the scale word needed let is = l / 16u; let sc_idx = sc_b_idx + is + quarter * 2u; - let sc_word = sc_idx / 4u; - let sc = bitcast(vec2( - src0[scale_idx + 96u + 2u * sc_word], - src0[scale_idx + 96u + 2u * sc_word + 1u] - )); - let sc_val = get_byte_i32(sc, sc_idx % 4u); - - let d = src0[scale_idx + 104u]; + let sc = load_src0_u32_at(block_byte_base + 192u + sc_idx); + let sc_val = get_byte_i32(sc, 0u); + + let d = load_src0_f16_at(block_byte_base + 208u); var q_val: f16; if (quarter == 0u) { diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl index 697126049062..aaf348eaced7 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl @@ -303,13 +303,7 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { #ifdef MUL_ACC_Q6_K const BLOCK_SIZE = 256u; -const F16_PER_BLOCK = 105u; - -fn load_u32_at(bbase: u32, byte_offset: u32) -> u32 { - let aligned = byte_offset & ~3u; - let idx = bbase + aligned / 2u; - return bitcast(vec2(src0[idx], src0[idx + 1u])); -} +const BLOCK_SIZE_BYTES = 210u; fn byte_of(v: u32, b: u32) -> u32 { return (v >> (b * 8u)) & 0xFFu; @@ -342,16 +336,15 @@ fn mul_acc(tig: u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { var local_sum = 0.0; for (var i = ix; i < nb; i += 2u) { - let bbase = (idx_base + k_block_start + i) * F16_PER_BLOCK; + let bbase = (idx_base + k_block_start + i) * BLOCK_SIZE_BYTES; - let d_raw = load_u32_at(bbase, 208u); - let d = f32(bitcast>(d_raw)[0]); + let d = f32(load_src0_f16_at(bbase + 208u)); - let ql1_u32 = load_u32_at(bbase, q_offset_l); - let ql2_u32 = load_u32_at(bbase, q_offset_l + 32u); - let qh_u32 = load_u32_at(bbase, 128u + q_offset_h); - let sc_u32_0 = load_u32_at(bbase, sc_base_byte); - let sc_u32_1 = load_u32_at(bbase, sc_base_byte + 4u); + let ql1_u32 = load_src0_u32_at(bbase + q_offset_l); + let ql2_u32 = load_src0_u32_at(bbase + q_offset_l + 32u); + let qh_u32 = load_src0_u32_at(bbase + 128u + q_offset_h); + let sc_u32_0 = load_src0_u32_at(bbase + sc_base_byte); + let sc_u32_1 = load_src0_u32_at(bbase + sc_base_byte + 4u); let sc0 = sbyte_of(sc_u32_0, sc_byte_pos); let sc2 = sbyte_of(sc_u32_0, sc_byte_pos + 2u); From b3aa3be8fe7a153cabfe963b32bda72a8c0f7f80 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Thu, 26 Mar 2026 08:59:23 -0700 Subject: [PATCH 03/16] Add timeout back to wait and remove synchronous set_tensor/memset_tensor --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 5e16f84ddd26..fdfed192bfaa 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -83,7 +83,7 @@ static inline void compute_2d_workgroups(uint32_t total_wg, uint32_t max_per_dim #define WEBGPU_NUM_PARAM_BUFS 96u #define WEBGPU_COMMAND_SUBMIT_BATCH_SIZE 32u -#define WEBGPU_WAIT_ANY_TIMEOUT_MS 0 +#define WEBGPU_WAIT_ANY_TIMEOUT_MS 100 // Maximum number of in-flight submissions per-thread, to avoid exhausting the // parameter buffer pool #define WEBGPU_MAX_INFLIGHT_SUBS_PER_THREAD (WEBGPU_NUM_PARAM_BUFS / WEBGPU_COMMAND_SUBMIT_BATCH_SIZE) @@ -514,7 +514,7 @@ static void ggml_backend_webgpu_wait(webgpu_global_context & ctx, bool blocking_wait = block || subs.size() >= WEBGPU_MAX_INFLIGHT_SUBS_PER_THREAD; while (blocking_wait) { - auto waitStatus = ctx->instance.WaitAny(1, &subs[0].submit_done, 0); + auto waitStatus = ctx->instance.WaitAny(1, &subs[0].submit_done, WEBGPU_WAIT_ANY_TIMEOUT_MS * 1e6); if (ggml_backend_webgpu_handle_wait_status(waitStatus, true)) { #ifdef GGML_WEBGPU_GPU_PROFILE ggml_backend_webgpu_wait_profile_futures(ctx, subs[0].profile_futures, true); @@ -735,7 +735,7 @@ static void ggml_backend_webgpu_buffer_memset(webgpu_global_context & ctx, ggml_backend_webgpu_build(ctx, ctx->memset_buf_pool, ctx->memset_pipelines[0], params, entries, wg_x); std::vector commands = { command }; std::vector sub = { ggml_backend_webgpu_submit(ctx, commands, ctx->memset_buf_pool) }; - ggml_backend_webgpu_wait(ctx, sub); + //ggml_backend_webgpu_wait(ctx, sub); } /** End WebGPU Actions */ @@ -2663,15 +2663,15 @@ static void ggml_backend_webgpu_buffer_set_tensor(ggml_backend_buffer_t buffer, total_offset + (size - remaining_size), remaining_size); } else { // wait for WriteBuffer to complete - buf_ctx->global_ctx->instance.WaitAny(buf_ctx->global_ctx->queue.OnSubmittedWorkDone( - wgpu::CallbackMode::AllowSpontaneous, - [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { - if (status != wgpu::QueueWorkDoneStatus::Success) { - GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", - std::string(message).c_str()); - } - }), - UINT64_MAX); +// buf_ctx->global_ctx->instance.WaitAny(buf_ctx->global_ctx->queue.OnSubmittedWorkDone( +// wgpu::CallbackMode::AllowSpontaneous, +// [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { +// if (status != wgpu::QueueWorkDoneStatus::Success) { +// GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", +// std::string(message).c_str()); +// } +// }), +// UINT64_MAX); } WEBGPU_CPU_PROFILE_TOTAL_END(set_tensor, buf_ctx->global_ctx); } From 67fe0897e82572d0d4e6c7fc3323899865e3b404 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Thu, 26 Mar 2026 09:41:37 -0700 Subject: [PATCH 04/16] move to unpackf16 for wider compatibility --- ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl | 3 ++- ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl | 3 ++- ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl index a721c0325a37..ce5af98270dc 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl @@ -95,7 +95,8 @@ fn load_v_u32_at(byte_offset: u32) -> u32 { } fn f16_from_u16(bits: u32) -> f16 { - return bitcast>(bits).x; + let packed = unpack2x16float(bits); + return f16(packed[0]); } #endif diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl index 4e58a892658c..1857d6112490 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl @@ -44,7 +44,8 @@ fn load_src0_u32_at(byte_offset: u32) -> u32 { } fn load_src0_f16_at(byte_offset: u32) -> f16 { - return bitcast>(load_src0_u16_at(byte_offset)).x; + let packed = unpack2x16float(load_src0_u16_at(byte_offset)); + return f16(packed[0]); } #endif // RAW_SRC0_BYTES diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl index aaf348eaced7..408210588ace 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl @@ -56,7 +56,8 @@ fn load_src0_u32_at(byte_offset: u32) -> u32 { } fn load_src0_f16_at(byte_offset: u32) -> f16 { - return bitcast>(load_src0_u16_at(byte_offset)).x; + let packed = unpack2x16float(load_src0_u16_at(byte_offset)); + return f16(packed[0]); } #endif From e85e8bcc0d5bda8a4c9b43d59fcbe5aba02ee875 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Thu, 26 Mar 2026 15:12:37 -0700 Subject: [PATCH 05/16] cleanup --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 38 ++----------------- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 12 ------ .../wgsl-shaders/common_decls.tmpl | 24 ++++++++++++ .../ggml-webgpu/wgsl-shaders/flash_attn.wgsl | 4 +- .../wgsl-shaders/mul_mat_decls.tmpl | 30 --------------- .../ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl | 30 --------------- 6 files changed, 30 insertions(+), 108 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index a3675b8ab058..277aa6375983 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -1124,20 +1124,8 @@ class ggml_webgpu_shader_lib { defines.push_back("BYTE_HELPERS"); defines.push_back("MUL_ACC_" + type_upper); - - if (context.src0->type == GGML_TYPE_Q4_0 || - context.src0->type == GGML_TYPE_Q4_1 || - context.src0->type == GGML_TYPE_Q5_0 || - context.src0->type == GGML_TYPE_Q5_1 || - context.src0->type == GGML_TYPE_Q8_0 || - context.src0->type == GGML_TYPE_Q8_1 || - context.src0->type == GGML_TYPE_Q6_K) { - defines.push_back("RAW_SRC0_BYTES"); - defines.push_back("SRC0_INNER_TYPE=u32"); - } else { - // For fast path we always dequantize from f16 inside the shader - defines.push_back("SRC0_INNER_TYPE=f16"); - } + defines.push_back("U32_DEQUANT_HELPERS"); + defines.push_back("SRC0_INNER_TYPE=u32"); break; } } @@ -1250,24 +1238,8 @@ class ggml_webgpu_shader_lib { defines.push_back("MUL_ACC_" + type_upper); defines.push_back("INIT_SRC0_SHMEM_" + type_upper); defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - - if (context.src0->type == GGML_TYPE_Q2_K || - context.src0->type == GGML_TYPE_Q3_K || - context.src0->type == GGML_TYPE_Q4_0 || - context.src0->type == GGML_TYPE_Q4_1 || - context.src0->type == GGML_TYPE_Q5_0 || - context.src0->type == GGML_TYPE_Q5_1 || - context.src0->type == GGML_TYPE_Q4_K || - context.src0->type == GGML_TYPE_Q5_K || - context.src0->type == GGML_TYPE_Q8_0 || - context.src0->type == GGML_TYPE_Q8_1 || - context.src0->type == GGML_TYPE_Q6_K) { - defines.push_back("RAW_SRC0_BYTES"); - defines.push_back("SRC0_INNER_TYPE=u32"); - } else { - // Use f16 inside the shader for quantized types - defines.push_back("SRC0_INNER_TYPE=f16"); - } + defines.push_back("U32_DEQUANT_HELPERS"); + defines.push_back("SRC0_INNER_TYPE=u32"); variant += std::string("_") + src0_name; break; @@ -1640,11 +1612,9 @@ class ggml_webgpu_shader_lib { break; case GGML_TYPE_Q4_0: defines.push_back("KV_Q4_0"); - defines.push_back("KV_RAW_U32"); break; case GGML_TYPE_Q8_0: defines.push_back("KV_Q8_0"); - defines.push_back("KV_RAW_U32"); break; default: GGML_ABORT("Unsupported KV type for flash attention shader"); diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index fdfed192bfaa..49a85d393fa0 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -735,7 +735,6 @@ static void ggml_backend_webgpu_buffer_memset(webgpu_global_context & ctx, ggml_backend_webgpu_build(ctx, ctx->memset_buf_pool, ctx->memset_pipelines[0], params, entries, wg_x); std::vector commands = { command }; std::vector sub = { ggml_backend_webgpu_submit(ctx, commands, ctx->memset_buf_pool) }; - //ggml_backend_webgpu_wait(ctx, sub); } /** End WebGPU Actions */ @@ -2661,17 +2660,6 @@ static void ggml_backend_webgpu_buffer_set_tensor(ggml_backend_buffer_t buffer, // memset the remaining bytes ggml_backend_webgpu_buffer_memset(buf_ctx->global_ctx, buf_ctx->buffer, val32, total_offset + (size - remaining_size), remaining_size); - } else { - // wait for WriteBuffer to complete -// buf_ctx->global_ctx->instance.WaitAny(buf_ctx->global_ctx->queue.OnSubmittedWorkDone( -// wgpu::CallbackMode::AllowSpontaneous, -// [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { -// if (status != wgpu::QueueWorkDoneStatus::Success) { -// GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", -// std::string(message).c_str()); -// } -// }), -// UINT64_MAX); } WEBGPU_CPU_PROFILE_TOTAL_END(set_tensor, buf_ctx->global_ctx); } diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl index 9a5b18ebc07b..feb0bca3f847 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl @@ -8,6 +8,30 @@ fn get_byte_i32(value: u32, index: u32) -> i32 { } #endif +#ifdef U32_DEQUANT_HELPERS +fn load_src0_u16_at(byte_offset: u32) -> u32 { + let word = src0[byte_offset / 4u]; + let shift = (byte_offset & 2u) * 8u; + return (word >> shift) & 0xFFFFu; +} + +fn load_src0_u32_at(byte_offset: u32) -> u32 { + let word_idx = byte_offset / 4u; + let shift = (byte_offset & 3u) * 8u; + let lo = src0[word_idx]; + if (shift == 0u) { + return lo; + } + let hi = src0[word_idx + 1u]; + return (lo >> shift) | (hi << (32u - shift)); +} + +fn load_src0_f16_at(byte_offset: u32) -> f16 { + let packed = unpack2x16float(load_src0_u16_at(byte_offset)); + return f16(packed[0]); +} +#endif + #ifdef Q4_0_T struct q4_0 { d: f16, diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl index ce5af98270dc..8b76cecba917 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl @@ -6,7 +6,7 @@ enable chromium_experimental_subgroup_matrix; #ifdef KV_F32 #define KV_TYPE f32 -#elif defined(KV_RAW_U32) +#elif defined(KV_Q4_0) || defined(KV_Q8_0) #define KV_TYPE u32 #else #define KV_TYPE f16 @@ -59,7 +59,7 @@ fn get_byte_i32(value: u32, index: u32) -> i32 { return bitcast(((value >> (index * 8)) & 0xFF) << 24) >> 24; } -#ifdef KV_RAW_U32 +#if defined(KV_Q4_0) || defined(KV_Q8_0) fn load_k_u16_at(byte_offset: u32) -> u32 { let word = K[byte_offset / 4u]; let shift = (byte_offset & 2u) * 8u; diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl index 1857d6112490..eb228537bad2 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl @@ -25,30 +25,6 @@ fn store_shmem(val: f16, idx: u32) { } #endif // SCALAR -#ifdef RAW_SRC0_BYTES -fn load_src0_u16_at(byte_offset: u32) -> u32 { - let word = src0[byte_offset / 4u]; - let shift = (byte_offset & 2u) * 8u; - return (word >> shift) & 0xFFFFu; -} - -fn load_src0_u32_at(byte_offset: u32) -> u32 { - let word_idx = byte_offset / 4u; - let shift = (byte_offset & 3u) * 8u; - let lo = src0[word_idx]; - if (shift == 0u) { - return lo; - } - let hi = src0[word_idx + 1u]; - return (lo >> shift) | (hi << (32u - shift)); -} - -fn load_src0_f16_at(byte_offset: u32) -> f16 { - let packed = unpack2x16float(load_src0_u16_at(byte_offset)); - return f16(packed[0]); -} -#endif // RAW_SRC0_BYTES - #ifdef INIT_SRC0_SHMEM_FLOAT fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { for (var elem_idx = thread_id * VEC_SIZE; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE * VEC_SIZE) { @@ -89,7 +65,6 @@ const BLOCK_SIZE_BYTES = 18u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; -const F16_PER_BLOCK = 9u; // 1 scale + 8x4 packed weights const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; @@ -131,7 +106,6 @@ const BLOCK_SIZE_BYTES = 20u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; -const F16_PER_BLOCK = 10u; // 1 scale + 8 packed weights + 1 mean const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; @@ -176,7 +150,6 @@ const BLOCK_SIZE_BYTES = 22u; // tile_k is defined as 32u, so blocks_k ends up being 1 always override BLOCKS_K = TILE_K / BLOCK_SIZE; const NQ = 16u; -const F16_PER_BLOCK = 11u; // 1 scale + 2 qh + 8 packed weights const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; // 16 / 4 = 4 f16s per thread, each thread should handle 4 f16s * 4 weights per = 16 weights @@ -231,7 +204,6 @@ const BLOCK_SIZE_BYTES = 24u; // tile_k is defined as 32u, so blocks_k ends up being 1 always override BLOCKS_K = TILE_K / BLOCK_SIZE; const NQ = 16u; -const F16_PER_BLOCK = 12u; // 1 scale + 2 qh + 8 packed weights + 1 mean const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; // 16 / 4 = 4 f16s per thread, each thread should handle 4 f16s * 4 weights per = 16 weights @@ -286,7 +258,6 @@ const BLOCK_SIZE_BYTES = 34u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; -const F16_PER_BLOCK = 17u; // 1 scale + 16 in array of weights const WEIGHTS_PER_F16 = 2u; // 2 8-bit weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; // 8 f16s per thread @@ -327,7 +298,6 @@ const BLOCK_SIZE_BYTES = 36u; // the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. override BLOCKS_K = TILE_K/BLOCK_SIZE; const NQ = 16u; -const F16_PER_BLOCK = 18u; // 1 scale + 1 mean + 8 32-bit values in array of weights const WEIGHTS_PER_F16 = 2u; // 2 8-bit weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; // 8 f16s per thread, 2 threads per block diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl index 408210588ace..6525f23bdfc4 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl @@ -37,30 +37,6 @@ fn store_val(group_base: u32) -> f32 { } #endif -#ifdef RAW_SRC0_BYTES -fn load_src0_u16_at(byte_offset: u32) -> u32 { - let word = src0[byte_offset / 4u]; - let shift = (byte_offset & 2u) * 8u; - return (word >> shift) & 0xFFFFu; -} - -fn load_src0_u32_at(byte_offset: u32) -> u32 { - let word_idx = byte_offset / 4u; - let shift = (byte_offset & 3u) * 8u; - let lo = src0[word_idx]; - if (shift == 0u) { - return lo; - } - let hi = src0[word_idx + 1u]; - return (lo >> shift) | (hi << (32u - shift)); -} - -fn load_src0_f16_at(byte_offset: u32) -> f16 { - let packed = unpack2x16float(load_src0_u16_at(byte_offset)); - return f16(packed[0]); -} -#endif - #ifdef MUL_ACC_FLOAT fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { var local_sum = 0.0; @@ -78,7 +54,6 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { const BLOCK_SIZE = 32; const BLOCK_SIZE_BYTES = 18u; const NQ = 16u; // number of weights per thread -const F16_PER_BLOCK = 9u; // 1 scale + 8x4 packed weights const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; @@ -112,7 +87,6 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { const BLOCK_SIZE = 32; const BLOCK_SIZE_BYTES = 20u; const NQ = 16u; // number of weights per thread -const F16_PER_BLOCK = 10u; const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; @@ -147,7 +121,6 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { const BLOCK_SIZE = 32; const BLOCK_SIZE_BYTES = 22u; const NQ = 16u; // number of weights per thread -const F16_PER_BLOCK = 11u; const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; @@ -192,7 +165,6 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { const BLOCK_SIZE = 32; const BLOCK_SIZE_BYTES = 24u; const NQ = 16u; // number of weights per thread -const F16_PER_BLOCK = 12u; const WEIGHTS_PER_F16 = 4u; // 4 weights per f16 const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; @@ -238,7 +210,6 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { const BLOCK_SIZE = 32; const BLOCK_SIZE_BYTES = 34u; const NQ = 16u; // number of weights per thread -const F16_PER_BLOCK = 17u; const WEIGHTS_PER_F16 = 2u; const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; @@ -272,7 +243,6 @@ fn mul_acc(tig:u32, tile_size: u32, idx_base: u32, k_outer: u32) -> f32 { const BLOCK_SIZE = 32; const BLOCK_SIZE_BYTES = 36u; const NQ = 16u; // number of weights per thread -const F16_PER_BLOCK = 18u; const WEIGHTS_PER_F16 = 2u; const F16_PER_THREAD = NQ / WEIGHTS_PER_F16; From 32ee70a2152004dd96b99e461cdb6e9c64908bc1 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Mon, 30 Mar 2026 10:15:26 -0700 Subject: [PATCH 06/16] Remove deadlock condition in free_bufs --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 49a85d393fa0..2cb551e483c0 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -171,6 +171,7 @@ struct webgpu_buf_pool { // Try growing the pool if no free buffers if (free.empty() && cur_pool_size < max_pool_size && should_grow) { cur_pool_size++; + lock.unlock(); // avoid deadlock between this lock and Dawn's internal locks when buffers are freed in callbacks wgpu::Buffer dev_buf; ggml_webgpu_create_buffer(device, dev_buf, buf_size, dev_buf_usage, "ggml_webgpu_dev_pool_buf"); From 1fc8b64a9274bee9e943136e4d453638d80b065c Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Tue, 31 Mar 2026 10:40:59 -0700 Subject: [PATCH 07/16] Start work on removing parameter buffer pools --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 258 +++++++++++++-------------- 1 file changed, 121 insertions(+), 137 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 2cb551e483c0..9813123fc3ba 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -81,12 +81,9 @@ static inline void compute_2d_workgroups(uint32_t total_wg, uint32_t max_per_dim /* Constants */ -#define WEBGPU_NUM_PARAM_BUFS 96u #define WEBGPU_COMMAND_SUBMIT_BATCH_SIZE 32u +#define WEBGPU_NUM_PARAM_SLOTS (WEBGPU_COMMAND_SUBMIT_BATCH_SIZE + 10) // a few extra for safety, since some operations may need multiple slots #define WEBGPU_WAIT_ANY_TIMEOUT_MS 100 -// Maximum number of in-flight submissions per-thread, to avoid exhausting the -// parameter buffer pool -#define WEBGPU_MAX_INFLIGHT_SUBS_PER_THREAD (WEBGPU_NUM_PARAM_BUFS / WEBGPU_COMMAND_SUBMIT_BATCH_SIZE) #define WEBGPU_PARAMS_BUF_SIZE_BYTES 128 // enough for 32 parameters #define WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES 4 #define WEBGPU_STORAGE_BUF_BINDING_MULT 4 // a storage buffer binding size must be a multiple of 4 @@ -122,87 +119,46 @@ static void ggml_webgpu_create_buffer(wgpu::Device & device, wgpu::BufferUsage usage, const char * label); -// Holds a pool of parameter buffers for WebGPU operations -struct webgpu_buf_pool { - std::vector free; - - // The pool must be synchronized because - // 1. The memset pool is shared globally by every ggml buffer, - // since allocating a pool per ggml buffer would consume too much memory. - // 2. For the per-thread buffer pools in webgpu_context, - // buffers are allocated and freed in Dawn callbacks, - // which can run on a different thread than the calling thread. - std::mutex mutex; - std::condition_variable cv; - size_t cur_pool_size; - size_t max_pool_size; - wgpu::Device device; - wgpu::BufferUsage dev_buf_usage; - size_t buf_size; - bool should_grow; - - void init(wgpu::Device device, - int num_bufs, - size_t buf_size, - wgpu::BufferUsage dev_buf_usage, - bool should_grow = false, - size_t max_pool_size = WEBGPU_NUM_PARAM_BUFS * 2) { - this->max_pool_size = max_pool_size; - this->cur_pool_size = num_bufs; - this->device = device; - this->dev_buf_usage = dev_buf_usage; - this->buf_size = buf_size; - this->should_grow = should_grow; - for (int i = 0; i < num_bufs; i++) { - wgpu::Buffer dev_buf; - ggml_webgpu_create_buffer(device, dev_buf, buf_size, dev_buf_usage, "ggml_webgpu_dev_pool_buf"); - free.push_back(dev_buf); - } - } - - wgpu::Buffer alloc_bufs() { - std::unique_lock lock(mutex); - if (!free.empty()) { - wgpu::Buffer buf = free.back(); - free.pop_back(); - return buf; +// Slot-based parameter arena for compute graph encoding. Each encoded kernel +// gets a unique uniform-buffer slice within the current batch, and the slot +// cursor is reset immediately after that batch is submitted. +struct webgpu_param_arena { + wgpu::Buffer buffer; + size_t slot_stride = 0; + size_t slot_size = 0; + uint32_t slot_count = 0; + uint32_t next_slot = 0; + + void init(wgpu::Device device, size_t slot_size, uint32_t slot_count, size_t alignment) { + this->slot_stride = ROUNDUP_POW2(slot_size, alignment); + this->slot_size = slot_size; + this->slot_count = slot_count; + this->next_slot = 0; + + ggml_webgpu_create_buffer(device, buffer, this->slot_stride * slot_count, + wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, + "ggml_webgpu_param_arena"); + } + + size_t alloc_slot(size_t size) { + GGML_ASSERT(size <= slot_size); + if (next_slot >= slot_count) { + GGML_ABORT("ggml_webgpu: parameter arena exhausted while encoding a batch"); } - // Try growing the pool if no free buffers - if (free.empty() && cur_pool_size < max_pool_size && should_grow) { - cur_pool_size++; - lock.unlock(); // avoid deadlock between this lock and Dawn's internal locks when buffers are freed in callbacks - wgpu::Buffer dev_buf; - ggml_webgpu_create_buffer(device, dev_buf, buf_size, dev_buf_usage, "ggml_webgpu_dev_pool_buf"); - - if (!dev_buf) { - GGML_ABORT("webgpu_buf_pool: failed to allocate buffers"); - } - return dev_buf; - } - cv.wait(lock, [this] { return !free.empty(); }); - wgpu::Buffer buf = free.back(); - free.pop_back(); - return buf; + return slot_stride * next_slot++; } - void free_bufs(std::vector bufs) { - std::lock_guard lock(mutex); - free.insert(free.end(), bufs.begin(), bufs.end()); - cv.notify_all(); - } + void reset() { next_slot = 0; } void cleanup() { - std::lock_guard lock(mutex); - for (auto & buf : free) { - if (buf) { - buf.Destroy(); - } + if (buffer) { + buffer.Destroy(); + buffer = nullptr; } - free.clear(); } - ~webgpu_buf_pool() { this->cleanup(); } + ~webgpu_param_arena() { this->cleanup(); } }; #ifdef GGML_WEBGPU_GPU_PROFILE @@ -272,7 +228,6 @@ struct webgpu_gpu_profile_buf_pool { struct webgpu_command { uint32_t num_kernels; wgpu::CommandBuffer commands; - std::vector params_bufs; #ifdef GGML_WEBGPU_GPU_PROFILE webgpu_gpu_profile_bufs timestamp_query_bufs; std::string pipeline_name; @@ -305,7 +260,7 @@ struct webgpu_global_context_struct { // Global mutex for pipeline and staging buffer, will be refactored to exclude pipeline caches. std::recursive_mutex mutex; - webgpu_buf_pool memset_buf_pool; + wgpu::Buffer memset_params_buf; std::map memset_pipelines; // variant or type index #ifdef GGML_WEBGPU_CPU_PROFILE @@ -332,6 +287,10 @@ struct webgpu_global_context_struct { this->get_tensor_staging_buf.Destroy(); this->get_tensor_staging_buf = nullptr; } + if (this->memset_params_buf) { + this->memset_params_buf.Destroy(); + this->memset_params_buf = nullptr; + } #ifdef GGML_WEBGPU_DEBUG if (this->debug_host_buf) { this->debug_host_buf.Destroy(); @@ -361,7 +320,7 @@ struct webgpu_context_struct { std::unique_ptr shader_lib; - webgpu_buf_pool param_buf_pool; + webgpu_param_arena param_arena; wgpu::Buffer set_rows_dev_error_buf; wgpu::Buffer set_rows_host_error_buf; @@ -513,7 +472,7 @@ static void ggml_backend_webgpu_wait(webgpu_global_context & ctx, return; } - bool blocking_wait = block || subs.size() >= WEBGPU_MAX_INFLIGHT_SUBS_PER_THREAD; + bool blocking_wait = block; while (blocking_wait) { auto waitStatus = ctx->instance.WaitAny(1, &subs[0].submit_done, WEBGPU_WAIT_ANY_TIMEOUT_MS * 1e6); if (ggml_backend_webgpu_handle_wait_status(waitStatus, true)) { @@ -522,7 +481,7 @@ static void ggml_backend_webgpu_wait(webgpu_global_context & ctx, #endif subs.erase(subs.begin()); } - blocking_wait = (block && !subs.empty()) || subs.size() >= WEBGPU_MAX_INFLIGHT_SUBS_PER_THREAD; + blocking_wait = block && !subs.empty(); } if (subs.empty()) { @@ -578,10 +537,8 @@ static void ggml_backend_webgpu_debug(webgpu_global_context & ctx) { #endif static webgpu_submission ggml_backend_webgpu_submit(webgpu_global_context & ctx, - std::vector & commands, - webgpu_buf_pool & param_buf_pool) { + std::vector & commands) { std::vector command_buffers; - std::vector params_bufs; webgpu_submission submission; #ifdef GGML_WEBGPU_GPU_PROFILE std::vector> pipeline_name_and_ts_bufs; @@ -589,18 +546,15 @@ static webgpu_submission ggml_backend_webgpu_submit(webgpu_global_context & for (const auto & command : commands) { command_buffers.push_back(command.commands); - params_bufs.insert(params_bufs.end(), command.params_bufs.begin(), command.params_bufs.end()); } ctx->queue.Submit(command_buffers.size(), command_buffers.data()); wgpu::Future p_f = ctx->queue.OnSubmittedWorkDone( wgpu::CallbackMode::AllowSpontaneous, - [¶m_buf_pool, params_bufs](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { + [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { if (status != wgpu::QueueWorkDoneStatus::Success) { GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", std::string(message).c_str()); } - // Free the staged buffers - param_buf_pool.free_bufs(params_bufs); }); submission.submit_done = { p_f }; @@ -631,7 +585,7 @@ static webgpu_submission ggml_backend_webgpu_submit(webgpu_global_context & static webgpu_command ggml_backend_webgpu_build_multi( webgpu_global_context & ctx, - webgpu_buf_pool & param_buf_pool, + webgpu_param_arena & param_arena, const std::vector & pipelines, const std::vector> & params_list, const std::vector> & bind_group_entries_list, @@ -640,16 +594,19 @@ static webgpu_command ggml_backend_webgpu_build_multi( GGML_ASSERT(pipelines.size() == bind_group_entries_list.size()); GGML_ASSERT(pipelines.size() == workgroups_list.size()); - std::vector params_bufs_list; std::vector bind_groups; + std::vector param_offsets; for (size_t i = 0; i < pipelines.size(); i++) { - wgpu::Buffer params_bufs = param_buf_pool.alloc_bufs(); + const size_t param_size = params_list[i].size() * sizeof(uint32_t); + const size_t param_offset = param_arena.alloc_slot(param_size); std::vector entries = bind_group_entries_list[i]; uint32_t params_binding_num = entries.size(); - entries.push_back( - { .binding = params_binding_num, .buffer = params_bufs, .offset = 0, .size = params_bufs.GetSize() }); + entries.push_back({ .binding = params_binding_num, + .buffer = param_arena.buffer, + .offset = param_offset, + .size = param_arena.slot_size }); wgpu::BindGroupDescriptor bind_group_desc; bind_group_desc.layout = pipelines[i].pipeline.GetBindGroupLayout(0); @@ -657,13 +614,13 @@ static webgpu_command ggml_backend_webgpu_build_multi( bind_group_desc.entries = entries.data(); bind_group_desc.label = pipelines[i].name.c_str(); bind_groups.push_back(ctx->device.CreateBindGroup(&bind_group_desc)); - - params_bufs_list.push_back(params_bufs); + param_offsets.push_back(param_offset); } wgpu::CommandEncoder encoder = ctx->device.CreateCommandEncoder(); - for (size_t i = 0; i < params_bufs_list.size(); i++) { - ctx->queue.WriteBuffer(params_bufs_list[i], 0, params_list[i].data(), params_list[i].size() * sizeof(uint32_t)); + for (size_t i = 0; i < param_offsets.size(); i++) { + ctx->queue.WriteBuffer(param_arena.buffer, param_offsets[i], params_list[i].data(), + params_list[i].size() * sizeof(uint32_t)); } #ifdef GGML_WEBGPU_GPU_PROFILE @@ -695,7 +652,6 @@ static webgpu_command ggml_backend_webgpu_build_multi( wgpu::CommandBuffer commands = encoder.Finish(); webgpu_command result = {}; result.commands = commands; - result.params_bufs = params_bufs_list; result.num_kernels = pipelines.size(); #ifdef GGML_WEBGPU_GPU_PROFILE result.timestamp_query_bufs = ts_bufs; @@ -706,13 +662,13 @@ static webgpu_command ggml_backend_webgpu_build_multi( } static webgpu_command ggml_backend_webgpu_build(webgpu_global_context & ctx, - webgpu_buf_pool & param_buf_pool, + webgpu_param_arena & param_arena, webgpu_pipeline & pipeline, std::vector params, std::vector bind_group_entries, uint32_t wg_x, uint32_t wg_y = 1) { - return ggml_backend_webgpu_build_multi(ctx, param_buf_pool, + return ggml_backend_webgpu_build_multi(ctx, param_arena, { pipeline }, @@ -732,10 +688,37 @@ static void ggml_backend_webgpu_buffer_memset(webgpu_global_context & ctx, size_t bytes_per_wg = WEBGPU_MAX_WG_SIZE * ctx->capabilities.memset_bytes_per_thread; uint32_t wg_x = CEIL_DIV(size + 3, bytes_per_wg); - webgpu_command command = - ggml_backend_webgpu_build(ctx, ctx->memset_buf_pool, ctx->memset_pipelines[0], params, entries, wg_x); - std::vector commands = { command }; - std::vector sub = { ggml_backend_webgpu_submit(ctx, commands, ctx->memset_buf_pool) }; + ctx->queue.WriteBuffer(ctx->memset_params_buf, 0, params.data(), params.size() * sizeof(uint32_t)); + + entries.push_back( + { .binding = 1, .buffer = ctx->memset_params_buf, .offset = 0, .size = WEBGPU_PARAMS_BUF_SIZE_BYTES }); + + wgpu::BindGroupDescriptor bind_group_desc; + bind_group_desc.layout = ctx->memset_pipelines[0].pipeline.GetBindGroupLayout(0); + bind_group_desc.entryCount = entries.size(); + bind_group_desc.entries = entries.data(); + bind_group_desc.label = ctx->memset_pipelines[0].name.c_str(); + wgpu::BindGroup bind_group = ctx->device.CreateBindGroup(&bind_group_desc); + + wgpu::CommandEncoder encoder = ctx->device.CreateCommandEncoder(); + wgpu::ComputePassEncoder pass = encoder.BeginComputePass(); + pass.SetPipeline(ctx->memset_pipelines[0].pipeline); + pass.SetBindGroup(0, bind_group); + pass.DispatchWorkgroups(wg_x, 1, 1); + pass.End(); + + wgpu::CommandBuffer command = encoder.Finish(); + std::vector commands = { command }; + ctx->queue.Submit(commands.size(), commands.data()); + ctx->instance.WaitAny(ctx->queue.OnSubmittedWorkDone( + wgpu::CallbackMode::AllowSpontaneous, + [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { + if (status != wgpu::QueueWorkDoneStatus::Success) { + GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", + std::string(message).c_str()); + } + }), + UINT64_MAX); } /** End WebGPU Actions */ @@ -876,7 +859,7 @@ static webgpu_command ggml_webgpu_cpy(webgpu_context & ctx, ggml_tensor * src, g }; uint32_t wg_x = CEIL_DIV(ne, WEBGPU_MAX_WG_SIZE); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, ctx->cpy_pipelines[src->type][dst->type], + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, ctx->cpy_pipelines[src->type][dst->type], params, entries, wg_x); } @@ -939,7 +922,7 @@ static webgpu_command ggml_webgpu_set(webgpu_context & ctx, ggml_tensor * src0, .size = ggml_webgpu_tensor_binding_size(ctx, dst) }); uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_pad(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { @@ -994,7 +977,7 @@ static webgpu_command ggml_webgpu_pad(webgpu_context & ctx, ggml_tensor * src, g }; uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_solve_tri(webgpu_context & ctx, @@ -1055,7 +1038,7 @@ static webgpu_command ggml_webgpu_solve_tri(webgpu_context & ctx, const uint32_t wg_x = CEIL_DIV((uint32_t) src1->ne[0], decisions->wg_size); const uint32_t wg_y = (uint32_t) (dst->ne[2] * dst->ne[3]); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x, wg_y); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, wg_y); } static webgpu_command ggml_webgpu_ssm_conv(webgpu_context & ctx, @@ -1111,7 +1094,7 @@ static webgpu_command ggml_webgpu_ssm_conv(webgpu_context & ctx, const uint32_t wg_x = CEIL_DIV((uint32_t) src0->ne[1], decisions->block_size); const uint32_t wg_y = token_tiles * (uint32_t) dst->ne[2]; - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x, wg_y); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, wg_y); } static webgpu_command ggml_webgpu_gated_delta_net(webgpu_context & ctx, @@ -1196,7 +1179,7 @@ static webgpu_command ggml_webgpu_gated_delta_net(webgpu_context & ctx, .size = ggml_webgpu_tensor_binding_size(ctx, dst) } }; - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, h, n_seqs); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, h, n_seqs); } static std::optional ggml_webgpu_set_rows(webgpu_context & ctx, @@ -1265,7 +1248,7 @@ static std::optional ggml_webgpu_set_rows(webgpu_context & ctx, threads = src->ne[0] * src->ne[1] * src->ne[2] * src->ne[3]; } uint32_t wg_x = CEIL_DIV(threads, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x, 1); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, 1); } // Workgroup size is a common constant @@ -1331,7 +1314,7 @@ static webgpu_command ggml_webgpu_get_rows(webgpu_context & ctx, uint32_t total_threads = float_parallel ? blocks_per_row * total_rows : total_rows; uint32_t wg_x = CEIL_DIV(total_threads, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_mul_mat(webgpu_context & ctx, @@ -1476,7 +1459,7 @@ static webgpu_command ggml_webgpu_mul_mat(webgpu_context & ctx, compute_2d_workgroups(total_wg, max_wg_per_dim, wg_x, wg_y); } - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x, wg_y); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, wg_y); } #ifndef __EMSCRIPTEN__ @@ -1584,7 +1567,7 @@ static webgpu_command ggml_webgpu_flash_attn(webgpu_context & ctx, uint32_t wg_per_head = CEIL_DIV(Q->ne[1], decisions->q_tile); uint32_t wg_x = wg_per_head * Q->ne[2] * Q->ne[3]; // wg per head * number of heads * number of batches - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } #endif @@ -1663,7 +1646,7 @@ static webgpu_command ggml_webgpu_unary_op(webgpu_context & ctx, ggml_tensor * s } uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_binary_op(webgpu_context & ctx, @@ -1765,7 +1748,7 @@ static webgpu_command ggml_webgpu_binary_op(webgpu_context & ctx, } uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_concat(webgpu_context & ctx, @@ -1821,7 +1804,7 @@ static webgpu_command ggml_webgpu_concat(webgpu_context & ctx, webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx); auto * decisions = static_cast(pipeline.context.get()); uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_repeat(webgpu_context & ctx, ggml_tensor * src0, ggml_tensor * dst) { @@ -1863,7 +1846,7 @@ static webgpu_command ggml_webgpu_repeat(webgpu_context & ctx, ggml_tensor * src webgpu_pipeline pipeline = ctx->shader_lib->get_repeat_pipeline(shader_lib_ctx); auto * decisions = static_cast(pipeline.context.get()); uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_row_norm(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { @@ -1906,7 +1889,7 @@ static webgpu_command ggml_webgpu_row_norm(webgpu_context & ctx, ggml_tensor * s }; webgpu_pipeline pipeline = ctx->shader_lib->get_row_norm_pipeline(shader_lib_ctx); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, ggml_nrows(src)); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, ggml_nrows(src)); } static webgpu_command ggml_webgpu_rope(webgpu_context & ctx, @@ -1998,7 +1981,7 @@ static webgpu_command ggml_webgpu_rope(webgpu_context & ctx, webgpu_pipeline pipeline = ctx->rope_pipelines[dst->type][has_freq_factor][inplace]; uint32_t wg_x = CEIL_DIV(ggml_nelements(dst), WEBGPU_MAX_WG_SIZE); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_glu(webgpu_context & ctx, ggml_tensor * src0, ggml_tensor * src1, ggml_tensor * dst) { @@ -2050,7 +2033,7 @@ static webgpu_command ggml_webgpu_glu(webgpu_context & ctx, ggml_tensor * src0, webgpu_pipeline pipeline = ctx->glu_pipelines[ggml_get_glu_op(dst)][dst->type][split]; uint32_t wg_x = CEIL_DIV(ggml_nelements(dst), WEBGPU_MAX_WG_SIZE); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_scale(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { @@ -2101,7 +2084,7 @@ static webgpu_command ggml_webgpu_scale(webgpu_context & ctx, ggml_tensor * src, } uint32_t wg_x = CEIL_DIV(ggml_nelements(dst), decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_soft_max(webgpu_context & ctx, @@ -2173,7 +2156,7 @@ static webgpu_command ggml_webgpu_soft_max(webgpu_context & ctx, .size = ggml_webgpu_tensor_binding_size(ctx, dst) }); } - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, ctx->soft_max_pipelines[mask_type][has_sink][inplace], params, entries, ggml_nrows(dst)); } @@ -2200,7 +2183,7 @@ static webgpu_command ggml_webgpu_argmax(webgpu_context & ctx, ggml_tensor * src webgpu_pipeline pipeline = ctx->shader_lib->get_argmax_pipeline(shader_lib_ctx); uint32_t wg_x = ggml_nelements(dst); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { @@ -2294,7 +2277,7 @@ static webgpu_command ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * sr workgroups_list.push_back({ wg_x_init, wg_y_init }); if (merge_passes == 0) { - return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_buf_pool, pipelines, params_list, + return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, pipelines, params_list, entries_list, workgroups_list); } @@ -2356,7 +2339,7 @@ static webgpu_command ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * sr in_is_tmp = !in_is_tmp; } - return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_buf_pool, pipelines, params_list, entries_list, + return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, pipelines, params_list, entries_list, workgroups_list); } @@ -2385,7 +2368,7 @@ static webgpu_command ggml_webgpu_cumsum(webgpu_context & ctx, ggml_tensor * src webgpu_pipeline pipeline = ctx->shader_lib->get_cumsum_pipeline(shader_lib_ctx); uint32_t wg_x = ggml_nrows(dst); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } static webgpu_command ggml_webgpu_sum_rows(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { @@ -2417,7 +2400,7 @@ static webgpu_command ggml_webgpu_sum_rows(webgpu_context & ctx, ggml_tensor * s webgpu_pipeline pipeline = ctx->shader_lib->get_sum_rows_pipeline(shader_lib_ctx); uint32_t wg_x = total_sum ? 1 : ggml_nrows(dst); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_buf_pool, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); } // Returns the encoded command, or std::nullopt if the operation is a no-op @@ -2538,15 +2521,17 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str if (num_batched_kernels >= WEBGPU_COMMAND_SUBMIT_BATCH_SIZE) { num_batched_kernels = 0; - subs.push_back(ggml_backend_webgpu_submit(ctx->global_ctx, commands, ctx->param_buf_pool)); + subs.push_back(ggml_backend_webgpu_submit(ctx->global_ctx, commands)); + ctx->param_arena.reset(); // Process events and check for completed submissions ctx->global_ctx->instance.ProcessEvents(); - ggml_backend_webgpu_wait(ctx->global_ctx, subs, false); + //ggml_backend_webgpu_wait(ctx->global_ctx, subs, false); commands.clear(); } } if (!commands.empty()) { - subs.push_back(ggml_backend_webgpu_submit(ctx->global_ctx, commands, ctx->param_buf_pool)); + subs.push_back(ggml_backend_webgpu_submit(ctx->global_ctx, commands)); + ctx->param_arena.reset(); commands.clear(); } @@ -3136,9 +3121,9 @@ static bool create_webgpu_device(ggml_backend_webgpu_reg_context * ctx) { GGML_ASSERT(ctx->webgpu_global_ctx->device != nullptr); ggml_webgpu_init_memset_pipeline(ctx->webgpu_global_ctx); - ctx->webgpu_global_ctx->memset_buf_pool.init(ctx->webgpu_global_ctx->device, 1, WEBGPU_PARAMS_BUF_SIZE_BYTES, - wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, - wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::MapWrite); + ggml_webgpu_create_buffer(ctx->webgpu_global_ctx->device, ctx->webgpu_global_ctx->memset_params_buf, + WEBGPU_PARAMS_BUF_SIZE_BYTES, + wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, "memset_params_buf"); ctx->webgpu_global_ctx->queue = ctx->webgpu_global_ctx->device.GetQueue(); #ifdef GGML_WEBGPU_GPU_PROFILE @@ -3162,9 +3147,8 @@ static webgpu_context initialize_webgpu_context(ggml_backend_dev_t dev) { webgpu_context webgpu_ctx = std::make_shared(); webgpu_ctx->global_ctx = dev_ctx->webgpu_global_ctx; webgpu_ctx->shader_lib = std::make_unique(dev_ctx->webgpu_global_ctx->device); - webgpu_ctx->param_buf_pool.init(webgpu_ctx->global_ctx->device, WEBGPU_NUM_PARAM_BUFS, WEBGPU_PARAMS_BUF_SIZE_BYTES, - wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, - wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::MapWrite, true); + webgpu_ctx->param_arena.init(webgpu_ctx->global_ctx->device, WEBGPU_PARAMS_BUF_SIZE_BYTES, WEBGPU_NUM_PARAM_SLOTS, + webgpu_ctx->global_ctx->capabilities.limits.minUniformBufferOffsetAlignment); ggml_webgpu_create_buffer(webgpu_ctx->global_ctx->device, webgpu_ctx->set_rows_dev_error_buf, WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES, wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc, "set_rows_dev_error_buf"); From 9592ed56af136a193ad794cc0023989fb487f25e Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Tue, 31 Mar 2026 11:41:15 -0700 Subject: [PATCH 08/16] Simplify and optimize further --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 174 ++++++++++----------------- 1 file changed, 65 insertions(+), 109 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 9813123fc3ba..e5b5d3d60368 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -81,8 +81,9 @@ static inline void compute_2d_workgroups(uint32_t total_wg, uint32_t max_per_dim /* Constants */ -#define WEBGPU_COMMAND_SUBMIT_BATCH_SIZE 32u -#define WEBGPU_NUM_PARAM_SLOTS (WEBGPU_COMMAND_SUBMIT_BATCH_SIZE + 10) // a few extra for safety, since some operations may need multiple slots +#define WEBGPU_COMMAND_SUBMIT_BATCH_SIZE 32u +#define WEBGPU_NUM_PARAM_SLOTS \ + (WEBGPU_COMMAND_SUBMIT_BATCH_SIZE + 10) // a few extra for safety, since some operations may need multiple slots #define WEBGPU_WAIT_ANY_TIMEOUT_MS 100 #define WEBGPU_PARAMS_BUF_SIZE_BYTES 128 // enough for 32 parameters #define WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES 4 @@ -136,8 +137,7 @@ struct webgpu_param_arena { this->next_slot = 0; ggml_webgpu_create_buffer(device, buffer, this->slot_stride * slot_count, - wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, - "ggml_webgpu_param_arena"); + wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, "ggml_webgpu_param_arena"); } size_t alloc_slot(size_t size) { @@ -226,8 +226,8 @@ struct webgpu_gpu_profile_buf_pool { #endif struct webgpu_command { - uint32_t num_kernels; - wgpu::CommandBuffer commands; + uint32_t num_kernels; + wgpu::CommandBuffer commands; #ifdef GGML_WEBGPU_GPU_PROFILE webgpu_gpu_profile_bufs timestamp_query_bufs; std::string pipeline_name; @@ -260,8 +260,8 @@ struct webgpu_global_context_struct { // Global mutex for pipeline and staging buffer, will be refactored to exclude pipeline caches. std::recursive_mutex mutex; - wgpu::Buffer memset_params_buf; - std::map memset_pipelines; // variant or type index + wgpu::Buffer memset_params_buf; + webgpu_pipeline memset_pipeline; #ifdef GGML_WEBGPU_CPU_PROFILE // Profiling: labeled CPU time in ms (total) @@ -306,13 +306,6 @@ struct webgpu_global_context_struct { typedef std::shared_ptr webgpu_global_context; -struct webgpu_submission { - wgpu::FutureWaitInfo submit_done; -#ifdef GGML_WEBGPU_GPU_PROFILE - std::vector profile_futures; -#endif -}; - // All the base objects needed to run operations on a WebGPU device struct webgpu_context_struct { // Points to global instances owned by ggml_backend_webgpu_reg_context @@ -321,8 +314,8 @@ struct webgpu_context_struct { std::unique_ptr shader_lib; webgpu_param_arena param_arena; - wgpu::Buffer set_rows_dev_error_buf; - wgpu::Buffer set_rows_host_error_buf; + wgpu::Buffer set_rows_dev_error_buf; + wgpu::Buffer set_rows_host_error_buf; std::map> cpy_pipelines; // src_type, dst_type @@ -464,45 +457,16 @@ static void ggml_backend_webgpu_wait_profile_futures(webgpu_global_context & } #endif -// Wait for the queue to finish processing all submitted work -static void ggml_backend_webgpu_wait(webgpu_global_context & ctx, - std::vector & subs, - bool block = true) { - if (subs.empty()) { - return; - } - - bool blocking_wait = block; - while (blocking_wait) { - auto waitStatus = ctx->instance.WaitAny(1, &subs[0].submit_done, WEBGPU_WAIT_ANY_TIMEOUT_MS * 1e6); - if (ggml_backend_webgpu_handle_wait_status(waitStatus, true)) { -#ifdef GGML_WEBGPU_GPU_PROFILE - ggml_backend_webgpu_wait_profile_futures(ctx, subs[0].profile_futures, true); -#endif - subs.erase(subs.begin()); - } - blocking_wait = block && !subs.empty(); - } - - if (subs.empty()) { - return; - } - - // Poll each submit future once and remove completed submissions. - for (auto sub = subs.begin(); sub != subs.end();) { - auto waitStatus = ctx->instance.WaitAny(1, &sub->submit_done, 0); - bool success = ggml_backend_webgpu_handle_wait_status(waitStatus, true); -#ifdef GGML_WEBGPU_GPU_PROFILE - ggml_backend_webgpu_wait_profile_futures(ctx, sub->profile_futures, false); - if (success && sub->profile_futures.empty()) { -#else - if (success) { -#endif - sub = subs.erase(sub); - } else { - ++sub; - } - } +static void ggml_backend_webgpu_wait_queue(webgpu_global_context & ctx) { + ctx->instance.WaitAny( + ctx->queue.OnSubmittedWorkDone(wgpu::CallbackMode::AllowSpontaneous, + [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { + if (status != wgpu::QueueWorkDoneStatus::Success) { + GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", + std::string(message).c_str()); + } + }), + UINT64_MAX); } static void ggml_backend_webgpu_map_buffer(webgpu_global_context & ctx, @@ -536,29 +500,19 @@ static void ggml_backend_webgpu_debug(webgpu_global_context & ctx) { } #endif -static webgpu_submission ggml_backend_webgpu_submit(webgpu_global_context & ctx, - std::vector & commands) { +static void ggml_backend_webgpu_submit(webgpu_global_context & ctx, std::vector & commands) { std::vector command_buffers; - webgpu_submission submission; -#ifdef GGML_WEBGPU_GPU_PROFILE - std::vector> pipeline_name_and_ts_bufs; -#endif - + command_buffers.reserve(commands.size()); for (const auto & command : commands) { command_buffers.push_back(command.commands); } ctx->queue.Submit(command_buffers.size(), command_buffers.data()); - - wgpu::Future p_f = ctx->queue.OnSubmittedWorkDone( - wgpu::CallbackMode::AllowSpontaneous, - [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { - if (status != wgpu::QueueWorkDoneStatus::Success) { - GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", std::string(message).c_str()); - } - }); - submission.submit_done = { p_f }; +} #ifdef GGML_WEBGPU_GPU_PROFILE +static void ggml_backend_webgpu_collect_profile_futures(webgpu_global_context & ctx, + const std::vector & commands, + std::vector & futures) { for (const auto & command : commands) { auto label = command.pipeline_name; auto ts_bufs = command.timestamp_query_bufs; @@ -577,11 +531,10 @@ static webgpu_submission ggml_backend_webgpu_submit(webgpu_global_context & // We can't unmap in here due to WebGPU reentrancy limitations. ctx->timestamp_query_buf_pool.free_bufs({ ts_bufs }); }); - submission.profile_futures.push_back({ f }); + futures.push_back({ f }); } -#endif - return submission; } +#endif static webgpu_command ggml_backend_webgpu_build_multi( webgpu_global_context & ctx, @@ -694,31 +647,22 @@ static void ggml_backend_webgpu_buffer_memset(webgpu_global_context & ctx, { .binding = 1, .buffer = ctx->memset_params_buf, .offset = 0, .size = WEBGPU_PARAMS_BUF_SIZE_BYTES }); wgpu::BindGroupDescriptor bind_group_desc; - bind_group_desc.layout = ctx->memset_pipelines[0].pipeline.GetBindGroupLayout(0); + bind_group_desc.layout = ctx->memset_pipeline.pipeline.GetBindGroupLayout(0); bind_group_desc.entryCount = entries.size(); bind_group_desc.entries = entries.data(); - bind_group_desc.label = ctx->memset_pipelines[0].name.c_str(); + bind_group_desc.label = ctx->memset_pipeline.name.c_str(); wgpu::BindGroup bind_group = ctx->device.CreateBindGroup(&bind_group_desc); - wgpu::CommandEncoder encoder = ctx->device.CreateCommandEncoder(); - wgpu::ComputePassEncoder pass = encoder.BeginComputePass(); - pass.SetPipeline(ctx->memset_pipelines[0].pipeline); + wgpu::CommandEncoder encoder = ctx->device.CreateCommandEncoder(); + wgpu::ComputePassEncoder pass = encoder.BeginComputePass(); + pass.SetPipeline(ctx->memset_pipeline.pipeline); pass.SetBindGroup(0, bind_group); pass.DispatchWorkgroups(wg_x, 1, 1); pass.End(); - wgpu::CommandBuffer command = encoder.Finish(); + wgpu::CommandBuffer command = encoder.Finish(); std::vector commands = { command }; ctx->queue.Submit(commands.size(), commands.data()); - ctx->instance.WaitAny(ctx->queue.OnSubmittedWorkDone( - wgpu::CallbackMode::AllowSpontaneous, - [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { - if (status != wgpu::QueueWorkDoneStatus::Success) { - GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", - std::string(message).c_str()); - } - }), - UINT64_MAX); } /** End WebGPU Actions */ @@ -2277,8 +2221,8 @@ static webgpu_command ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * sr workgroups_list.push_back({ wg_x_init, wg_y_init }); if (merge_passes == 0) { - return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, pipelines, params_list, - entries_list, workgroups_list); + return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, pipelines, params_list, entries_list, + workgroups_list); } bool in_is_tmp = start_in_tmp; @@ -2505,10 +2449,12 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str WEBGPU_CPU_PROFILE_TOTAL_START(graph_compute); - std::vector commands; - std::vector subs; - uint32_t num_batched_kernels = 0; - bool contains_set_rows = false; + std::vector commands; +#ifdef GGML_WEBGPU_GPU_PROFILE + std::vector profile_futures; +#endif + uint32_t num_batched_kernels = 0; + bool contains_set_rows = false; for (int i = 0; i < cgraph->n_nodes; i++) { if (cgraph->nodes[i]->op == GGML_OP_SET_ROWS) { @@ -2521,16 +2467,19 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str if (num_batched_kernels >= WEBGPU_COMMAND_SUBMIT_BATCH_SIZE) { num_batched_kernels = 0; - subs.push_back(ggml_backend_webgpu_submit(ctx->global_ctx, commands)); + ggml_backend_webgpu_submit(ctx->global_ctx, commands); +#ifdef GGML_WEBGPU_GPU_PROFILE + ggml_backend_webgpu_collect_profile_futures(ctx->global_ctx, commands, profile_futures); +#endif ctx->param_arena.reset(); - // Process events and check for completed submissions - ctx->global_ctx->instance.ProcessEvents(); - //ggml_backend_webgpu_wait(ctx->global_ctx, subs, false); commands.clear(); } } if (!commands.empty()) { - subs.push_back(ggml_backend_webgpu_submit(ctx->global_ctx, commands)); + ggml_backend_webgpu_submit(ctx->global_ctx, commands); +#ifdef GGML_WEBGPU_GPU_PROFILE + ggml_backend_webgpu_collect_profile_futures(ctx->global_ctx, commands, profile_futures); +#endif ctx->param_arena.reset(); commands.clear(); } @@ -2542,6 +2491,11 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str ctx->set_rows_host_error_buf.GetSize()); wgpu::CommandBuffer set_rows_commands = encoder.Finish(); ctx->global_ctx->queue.Submit(1, &set_rows_commands); + } + + ggml_backend_webgpu_wait_queue(ctx->global_ctx); + + if (contains_set_rows) { ggml_backend_webgpu_map_buffer(ctx->global_ctx, ctx->set_rows_host_error_buf, wgpu::MapMode::Read, 0, ctx->set_rows_host_error_buf.GetSize()); const uint32_t * error_data = (const uint32_t *) ctx->set_rows_host_error_buf.GetConstMappedRange(); @@ -2551,7 +2505,9 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str ctx->set_rows_host_error_buf.Unmap(); } - ggml_backend_webgpu_wait(ctx->global_ctx, subs); +#ifdef GGML_WEBGPU_GPU_PROFILE + ggml_backend_webgpu_wait_profile_futures(ctx->global_ctx, profile_futures, true); +#endif WEBGPU_CPU_PROFILE_TOTAL_END(graph_compute, ctx->global_ctx); return GGML_STATUS_SUCCESS; } @@ -2852,11 +2808,11 @@ static void ggml_webgpu_init_memset_pipeline(webgpu_global_context & ctx) { ctx->capabilities.memset_bytes_per_thread = CEIL_DIV(ctx->capabilities.limits.maxStorageBufferBindingSize, max_threads); std::vector constants(2); - constants[0].key = "wg_size"; - constants[0].value = WEBGPU_MAX_WG_SIZE; - constants[1].key = "bytes_per_thread"; - constants[1].value = ctx->capabilities.memset_bytes_per_thread; - ctx->memset_pipelines[0] = ggml_webgpu_create_pipeline(ctx->device, wgsl_memset, "memset", constants); + constants[0].key = "wg_size"; + constants[0].value = WEBGPU_MAX_WG_SIZE; + constants[1].key = "bytes_per_thread"; + constants[1].value = ctx->capabilities.memset_bytes_per_thread; + ctx->memset_pipeline = ggml_webgpu_create_pipeline(ctx->device, wgsl_memset, "memset", constants); } static void ggml_webgpu_init_cpy_pipeline(webgpu_context & webgpu_ctx) { @@ -3122,8 +3078,8 @@ static bool create_webgpu_device(ggml_backend_webgpu_reg_context * ctx) { ggml_webgpu_init_memset_pipeline(ctx->webgpu_global_ctx); ggml_webgpu_create_buffer(ctx->webgpu_global_ctx->device, ctx->webgpu_global_ctx->memset_params_buf, - WEBGPU_PARAMS_BUF_SIZE_BYTES, - wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, "memset_params_buf"); + WEBGPU_PARAMS_BUF_SIZE_BYTES, wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Uniform, + "memset_params_buf"); ctx->webgpu_global_ctx->queue = ctx->webgpu_global_ctx->device.GetQueue(); #ifdef GGML_WEBGPU_GPU_PROFILE From d307d47c12329f3db10a97af478c4b211feb4e14 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Wed, 1 Apr 2026 11:41:19 -0700 Subject: [PATCH 09/16] simplify profile futures --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 79 ++++------------------------ 1 file changed, 9 insertions(+), 70 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 396811f26be1..7ff5d66241a9 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -131,7 +131,7 @@ struct webgpu_param_arena { uint32_t next_slot = 0; void init(wgpu::Device device, size_t slot_size, uint32_t slot_count, size_t alignment) { - this->slot_stride = ROUNDUP_POW2(slot_size, alignment); + this->slot_stride = alignment; this->slot_size = slot_size; this->slot_count = slot_count; this->next_slot = 0; @@ -400,81 +400,20 @@ static void ggml_webgpu_create_buffer(wgpu::Device & device, /** WebGPU Actions */ -static bool ggml_backend_webgpu_handle_wait_status(wgpu::WaitStatus status, bool allow_timeout = false) { - switch (status) { - case wgpu::WaitStatus::Success: - return true; - case wgpu::WaitStatus::TimedOut: - if (allow_timeout) { - return false; - } - GGML_LOG_ERROR("ggml_webgpu: WaitAny timed out unexpectedly\n"); - return false; - case wgpu::WaitStatus::Error: - GGML_LOG_ERROR("ggml_webgpu: WaitAny returned an error\n"); - return false; - default: - GGML_LOG_ERROR("ggml_webgpu: WaitAny returned an unknown status\n"); - return false; - } -} - #ifdef GGML_WEBGPU_GPU_PROFILE -static void ggml_backend_webgpu_erase_completed_futures(std::vector & futures) { - futures.erase(std::remove_if(futures.begin(), futures.end(), - [](const wgpu::FutureWaitInfo & info) { return info.completed; }), - futures.end()); -} - -static size_t ggml_backend_webgpu_wait_profile_futures_pass(webgpu_global_context & ctx, - std::vector & futures, - uint64_t timeout_ms, - bool allow_timeout) { - constexpr size_t max_futures_per_wait = 64; - - size_t completed_before = futures.size(); - for (size_t offset = 0; offset < futures.size();) { - const size_t count = std::min(max_futures_per_wait, futures.size() - offset); - auto wait_status = ctx->instance.WaitAny(count, futures.data() + offset, timeout_ms); - - if (!ggml_backend_webgpu_handle_wait_status(wait_status, allow_timeout)) { - break; - } - - const size_t size_before_erase = futures.size(); - ggml_backend_webgpu_erase_completed_futures(futures); - if (futures.empty()) { - break; - } - - if (futures.size() != size_before_erase) { - offset = 0; - continue; - } - - offset += max_futures_per_wait; - } - - return completed_before - futures.size(); -} - static void ggml_backend_webgpu_wait_profile_futures(webgpu_global_context & ctx, - std::vector & futures, - bool block) { + std::vector & futures) { if (futures.empty()) { return; } - if (block) { - while (!futures.empty()) { - if (ggml_backend_webgpu_wait_profile_futures_pass(ctx, futures, 0, true) > 0) { - continue; - } + constexpr size_t max_futures_per_wait = 64; - ggml_backend_webgpu_wait_profile_futures_pass(ctx, futures, UINT64_MAX, false); - } - } else { - ggml_backend_webgpu_wait_profile_futures_pass(ctx, futures, 0, true); + while (!futures.empty()) { + ctx->instance.WaitAny(std::min(max_futures_per_wait, futures.size()), futures.data(), UINT64_MAX); + futures.erase(std::remove_if(futures.begin(), futures.end(), + [](const wgpu::FutureWaitInfo & info) { return info.completed; }), + futures.end()); } } #endif @@ -2568,7 +2507,7 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str } #ifdef GGML_WEBGPU_GPU_PROFILE - ggml_backend_webgpu_wait_profile_futures(ctx->global_ctx, profile_futures, true); + ggml_backend_webgpu_wait_profile_futures(ctx->global_ctx, profile_futures); #endif WEBGPU_CPU_PROFILE_TOTAL_END(graph_compute, ctx->global_ctx); return GGML_STATUS_SUCCESS; From a2c1d9102848521882c64f1b174acd91bb021298 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Wed, 1 Apr 2026 16:08:11 -0700 Subject: [PATCH 10/16] Fix stride --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 7ff5d66241a9..e662e4fa2d01 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -131,7 +131,7 @@ struct webgpu_param_arena { uint32_t next_slot = 0; void init(wgpu::Device device, size_t slot_size, uint32_t slot_count, size_t alignment) { - this->slot_stride = alignment; + this->slot_stride = ROUNDUP_POW2(slot_size, alignment); this->slot_size = slot_size; this->slot_count = slot_count; this->next_slot = 0; From e8ea0441e21d546c37315e9740a4a7dbc1b06839 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Thu, 2 Apr 2026 11:29:31 -0700 Subject: [PATCH 11/16] Try using a single command buffer per batch --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 298 +++++++++++++-------------- 1 file changed, 149 insertions(+), 149 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index e662e4fa2d01..ff22911ae177 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -225,9 +225,8 @@ struct webgpu_gpu_profile_buf_pool { }; #endif -struct webgpu_command { - uint32_t num_kernels; - wgpu::CommandBuffer commands; +struct webgpu_encoded_op { + uint32_t num_kernels = 0; #ifdef GGML_WEBGPU_GPU_PROFILE webgpu_gpu_profile_bufs timestamp_query_bufs; std::string pipeline_name; @@ -461,15 +460,6 @@ static void ggml_backend_webgpu_debug(webgpu_global_context & ctx) { } #endif -static void ggml_backend_webgpu_submit(webgpu_global_context & ctx, std::vector & commands) { - std::vector command_buffers; - command_buffers.reserve(commands.size()); - for (const auto & command : commands) { - command_buffers.push_back(command.commands); - } - ctx->queue.Submit(command_buffers.size(), command_buffers.data()); -} - #ifdef GGML_WEBGPU_GPU_PROFILE static void ggml_backend_webgpu_collect_profile_futures(webgpu_global_context & ctx, const std::vector & commands, @@ -497,9 +487,10 @@ static void ggml_backend_webgpu_collect_profile_futures(webgpu_global_context & } #endif -static webgpu_command ggml_backend_webgpu_build_multi( +static webgpu_encoded_op ggml_backend_webgpu_build_multi( webgpu_global_context & ctx, webgpu_param_arena & param_arena, + wgpu::CommandEncoder & encoder, const std::vector & pipelines, const std::vector> & params_list, const std::vector> & bind_group_entries_list, @@ -508,8 +499,10 @@ static webgpu_command ggml_backend_webgpu_build_multi( GGML_ASSERT(pipelines.size() == bind_group_entries_list.size()); GGML_ASSERT(pipelines.size() == workgroups_list.size()); + webgpu_encoded_op result = {}; std::vector bind_groups; std::vector param_offsets; + result.num_kernels = pipelines.size(); for (size_t i = 0; i < pipelines.size(); i++) { const size_t param_size = params_list[i].size() * sizeof(uint32_t); @@ -531,7 +524,6 @@ static webgpu_command ggml_backend_webgpu_build_multi( param_offsets.push_back(param_offset); } - wgpu::CommandEncoder encoder = ctx->device.CreateCommandEncoder(); for (size_t i = 0; i < param_offsets.size(); i++) { ctx->queue.WriteBuffer(param_arena.buffer, param_offsets[i], params_list[i].data(), params_list[i].size() * sizeof(uint32_t)); @@ -561,28 +553,21 @@ static webgpu_command ggml_backend_webgpu_build_multi( #ifdef GGML_WEBGPU_GPU_PROFILE encoder.ResolveQuerySet(ts_bufs.query_set, 0, 2, ts_bufs.dev_buf, 0); encoder.CopyBufferToBuffer(ts_bufs.dev_buf, 0, ts_bufs.host_buf, 0, ts_bufs.host_buf.GetSize()); -#endif - - wgpu::CommandBuffer commands = encoder.Finish(); - webgpu_command result = {}; - result.commands = commands; - result.num_kernels = pipelines.size(); -#ifdef GGML_WEBGPU_GPU_PROFILE result.timestamp_query_bufs = ts_bufs; - // TODO: handle multiple pipeline names result.pipeline_name = pipelines.front().name; #endif return result; } -static webgpu_command ggml_backend_webgpu_build(webgpu_global_context & ctx, - webgpu_param_arena & param_arena, - webgpu_pipeline & pipeline, - std::vector params, - std::vector bind_group_entries, - uint32_t wg_x, - uint32_t wg_y = 1) { - return ggml_backend_webgpu_build_multi(ctx, param_arena, +static webgpu_encoded_op ggml_backend_webgpu_build(webgpu_global_context & ctx, + webgpu_param_arena & param_arena, + wgpu::CommandEncoder & encoder, + webgpu_pipeline & pipeline, + std::vector params, + std::vector bind_group_entries, + uint32_t wg_x, + uint32_t wg_y = 1) { + return ggml_backend_webgpu_build_multi(ctx, param_arena, encoder, { pipeline }, @@ -736,7 +721,7 @@ static binary_overlap_flags ggml_webgpu_detect_binary_overlap(ggml_tensor * src0 return flags; } -static webgpu_command ggml_webgpu_cpy(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_cpy(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src, .dst = dst, @@ -774,10 +759,10 @@ static webgpu_command ggml_webgpu_cpy(webgpu_context & ctx, ggml_tensor * src, g }; uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_set(webgpu_context & ctx, ggml_tensor * src0, ggml_tensor * src1, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_set(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src0, ggml_tensor * src1, ggml_tensor * dst) { const bool inplace = ggml_webgpu_tensor_equal(src0, dst); ggml_webgpu_shader_lib_context shader_lib_ctx = { @@ -836,10 +821,10 @@ static webgpu_command ggml_webgpu_set(webgpu_context & ctx, ggml_tensor * src0, .size = ggml_webgpu_tensor_binding_size(ctx, dst) }); uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_pad(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_pad(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src, .dst = dst, .max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup }; @@ -891,13 +876,14 @@ static webgpu_command ggml_webgpu_pad(webgpu_context & ctx, ggml_tensor * src, g }; uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_solve_tri(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_solve_tri(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src0, .src1 = src1, @@ -952,13 +938,14 @@ static webgpu_command ggml_webgpu_solve_tri(webgpu_context & ctx, const uint32_t wg_x = CEIL_DIV((uint32_t) src1->ne[0], decisions->wg_size); const uint32_t wg_y = (uint32_t) (dst->ne[2] * dst->ne[3]); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, wg_y); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x, wg_y); } -static webgpu_command ggml_webgpu_ssm_conv(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_ssm_conv(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src0, .src1 = src1, @@ -1008,17 +995,18 @@ static webgpu_command ggml_webgpu_ssm_conv(webgpu_context & ctx, const uint32_t wg_x = CEIL_DIV((uint32_t) src0->ne[1], decisions->block_size); const uint32_t wg_y = token_tiles * (uint32_t) dst->ne[2]; - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, wg_y); -} - -static webgpu_command ggml_webgpu_gated_delta_net(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * src2, - ggml_tensor * src3, - ggml_tensor * src4, - ggml_tensor * src5, - ggml_tensor * dst) { + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x, wg_y); +} + +static webgpu_encoded_op ggml_webgpu_gated_delta_net(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * src2, + ggml_tensor * src3, + ggml_tensor * src4, + ggml_tensor * src5, + ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src0, .src1 = src1, @@ -1093,13 +1081,14 @@ static webgpu_command ggml_webgpu_gated_delta_net(webgpu_context & ctx, .size = ggml_webgpu_tensor_binding_size(ctx, dst) } }; - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, h, n_seqs); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, h, n_seqs); } -static std::optional ggml_webgpu_set_rows(webgpu_context & ctx, - ggml_tensor * src, - ggml_tensor * idx, - ggml_tensor * dst) { +static std::optional ggml_webgpu_set_rows(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src, + ggml_tensor * idx, + ggml_tensor * dst) { // For set rows specifically, we need to check if src and idx are empty // tensors. if (ggml_is_empty(src) || ggml_is_empty(idx)) { @@ -1162,7 +1151,7 @@ static std::optional ggml_webgpu_set_rows(webgpu_context & ctx, threads = src->ne[0] * src->ne[1] * src->ne[2] * src->ne[3]; } uint32_t wg_x = CEIL_DIV(threads, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, 1); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x, 1); } // Workgroup size is a common constant @@ -1173,10 +1162,11 @@ static std::vector ggml_webgpu_wg_size_entry(uint32_t wg_si return constants; } -static webgpu_command ggml_webgpu_get_rows(webgpu_context & ctx, - ggml_tensor * src, - ggml_tensor * idx, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_get_rows(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src, + ggml_tensor * idx, + ggml_tensor * dst) { const bool float_parallel = src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16 || src->type == GGML_TYPE_I32; ggml_webgpu_shader_lib_context shader_lib_ctx = { @@ -1228,13 +1218,14 @@ static webgpu_command ggml_webgpu_get_rows(webgpu_context & ctx, uint32_t total_threads = float_parallel ? blocks_per_row * total_rows : total_rows; uint32_t wg_x = CEIL_DIV(total_threads, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_mul_mat(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_mul_mat(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * dst) { // Determine if this is a mat-vec operation bool is_vec = (dst->ne[1] == 1); @@ -1373,17 +1364,18 @@ static webgpu_command ggml_webgpu_mul_mat(webgpu_context & ctx, compute_2d_workgroups(total_wg, max_wg_per_dim, wg_x, wg_y); } - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x, wg_y); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x, wg_y); } #ifndef __EMSCRIPTEN__ -static webgpu_command ggml_webgpu_flash_attn(webgpu_context & ctx, - ggml_tensor * Q, - ggml_tensor * K, - ggml_tensor * V, - ggml_tensor * mask, - ggml_tensor * sinks, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_flash_attn(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * Q, + ggml_tensor * K, + ggml_tensor * V, + ggml_tensor * mask, + ggml_tensor * sinks, + ggml_tensor * dst) { float scale = *(float *) dst->op_params; float max_bias; memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); @@ -1481,11 +1473,11 @@ static webgpu_command ggml_webgpu_flash_attn(webgpu_context & ctx, uint32_t wg_per_head = CEIL_DIV(Q->ne[1], decisions->q_tile); uint32_t wg_x = wg_per_head * Q->ne[2] * Q->ne[3]; // wg per head * number of heads * number of batches - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } #endif -static webgpu_command ggml_webgpu_unary_op(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_unary_op(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { bool is_unary = dst->op == GGML_OP_UNARY; bool inplace = ggml_webgpu_tensor_equal(src, dst) || (dst->op == GGML_OP_FILL); @@ -1560,13 +1552,14 @@ static webgpu_command ggml_webgpu_unary_op(webgpu_context & ctx, ggml_tensor * s } uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_binary_op(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_binary_op(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * dst) { binary_overlap_flags flags = ggml_webgpu_detect_binary_overlap(src0, src1, dst); ggml_webgpu_shader_lib_context shader_lib_ctx = { @@ -1662,13 +1655,14 @@ static webgpu_command ggml_webgpu_binary_op(webgpu_context & ctx, } uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_concat(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_concat(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * dst) { uint32_t ne = (uint32_t) ggml_nelements(dst); uint32_t dim = (uint32_t) dst->op_params[0]; @@ -1718,10 +1712,10 @@ static webgpu_command ggml_webgpu_concat(webgpu_context & ctx, webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx); auto * decisions = static_cast(pipeline.context.get()); uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_repeat(webgpu_context & ctx, ggml_tensor * src0, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_repeat(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src0, ggml_tensor * dst) { uint32_t ne = (uint32_t) ggml_nelements(dst); std::vector params = { ne, @@ -1760,10 +1754,10 @@ static webgpu_command ggml_webgpu_repeat(webgpu_context & ctx, ggml_tensor * src webgpu_pipeline pipeline = ctx->shader_lib->get_repeat_pipeline(shader_lib_ctx); auto * decisions = static_cast(pipeline.context.get()); uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_row_norm(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_row_norm(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { bool inplace = ggml_webgpu_tensor_equal(src, dst); std::vector params = { @@ -1803,14 +1797,15 @@ static webgpu_command ggml_webgpu_row_norm(webgpu_context & ctx, ggml_tensor * s }; webgpu_pipeline pipeline = ctx->shader_lib->get_row_norm_pipeline(shader_lib_ctx); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, ggml_nrows(src)); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, ggml_nrows(src)); } -static webgpu_command ggml_webgpu_rope(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * src2, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_rope(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * src2, + ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src0, .src1 = src1, @@ -1907,10 +1902,10 @@ static webgpu_command ggml_webgpu_rope(webgpu_context & ctx, } uint32_t wg_x = CEIL_DIV(ggml_nelements(dst), decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_glu(webgpu_context & ctx, ggml_tensor * src0, ggml_tensor * src1, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_glu(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src0, ggml_tensor * src1, ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src0, .src1 = src1, @@ -1969,10 +1964,10 @@ static webgpu_command ggml_webgpu_glu(webgpu_context & ctx, ggml_tensor * src0, .size = ggml_webgpu_tensor_binding_size(ctx, dst) }); uint32_t wg_x = CEIL_DIV(ggml_nelements(dst), decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_scale(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_scale(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { bool inplace = ggml_webgpu_tensor_equal(src, dst); ggml_webgpu_shader_lib_context shader_lib_ctx = { @@ -2020,14 +2015,15 @@ static webgpu_command ggml_webgpu_scale(webgpu_context & ctx, ggml_tensor * src, } uint32_t wg_x = CEIL_DIV(ggml_nelements(dst), decisions->wg_size); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_soft_max(webgpu_context & ctx, - ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * src2, - ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_soft_max(webgpu_context & ctx, + wgpu::CommandEncoder & encoder, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * src2, + ggml_tensor * dst) { ggml_webgpu_shader_lib_context shader_lib_ctx = { .src0 = src0, .src1 = src1, @@ -2103,10 +2099,10 @@ static webgpu_command ggml_webgpu_soft_max(webgpu_context & ctx, .size = ggml_webgpu_tensor_binding_size(ctx, dst) }); } - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, ggml_nrows(dst)); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, ggml_nrows(dst)); } -static webgpu_command ggml_webgpu_argmax(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_argmax(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { std::vector params = { (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)), (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), (uint32_t) src->ne[0] }; @@ -2128,10 +2124,10 @@ static webgpu_command ggml_webgpu_argmax(webgpu_context & ctx, ggml_tensor * src webgpu_pipeline pipeline = ctx->shader_lib->get_argmax_pipeline(shader_lib_ctx); uint32_t wg_x = ggml_nelements(dst); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_argsort(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { bool is_top_k = dst->op == GGML_OP_TOP_K; ggml_webgpu_shader_lib_context shader_lib_ctx = { @@ -2222,7 +2218,7 @@ static webgpu_command ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * sr workgroups_list.push_back({ wg_x_init, wg_y_init }); if (merge_passes == 0) { - return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, pipelines, params_list, entries_list, + return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, encoder, pipelines, params_list, entries_list, workgroups_list); } @@ -2284,11 +2280,11 @@ static webgpu_command ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * sr in_is_tmp = !in_is_tmp; } - return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, pipelines, params_list, entries_list, + return ggml_backend_webgpu_build_multi(ctx->global_ctx, ctx->param_arena, encoder, pipelines, params_list, entries_list, workgroups_list); } -static webgpu_command ggml_webgpu_cumsum(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_cumsum(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { std::vector params = { (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)), (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), (uint32_t) src->ne[0] }; @@ -2313,10 +2309,10 @@ static webgpu_command ggml_webgpu_cumsum(webgpu_context & ctx, ggml_tensor * src webgpu_pipeline pipeline = ctx->shader_lib->get_cumsum_pipeline(shader_lib_ctx); uint32_t wg_x = ggml_nrows(dst); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } -static webgpu_command ggml_webgpu_sum_rows(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) { +static webgpu_encoded_op ggml_webgpu_sum_rows(webgpu_context & ctx, wgpu::CommandEncoder & encoder, ggml_tensor * src, ggml_tensor * dst) { bool total_sum = dst->op == GGML_OP_SUM; std::vector params = { (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)), (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), @@ -2345,11 +2341,11 @@ static webgpu_command ggml_webgpu_sum_rows(webgpu_context & ctx, ggml_tensor * s webgpu_pipeline pipeline = ctx->shader_lib->get_sum_rows_pipeline(shader_lib_ctx); uint32_t wg_x = total_sum ? 1 : ggml_nrows(dst); - return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, pipeline, params, entries, wg_x); + return ggml_backend_webgpu_build(ctx->global_ctx, ctx->param_arena, encoder, pipeline, params, entries, wg_x); } // Returns the encoded command, or std::nullopt if the operation is a no-op -static std::optional ggml_webgpu_encode_node(webgpu_context ctx, ggml_tensor * node) { +static std::optional ggml_webgpu_encode_node(webgpu_context ctx, wgpu::CommandEncoder & encoder, ggml_tensor * node) { if (ggml_is_empty(node)) { return std::nullopt; } @@ -2372,18 +2368,18 @@ static std::optional ggml_webgpu_encode_node(webgpu_context ctx, return std::nullopt; case GGML_OP_CPY: case GGML_OP_CONT: - return ggml_webgpu_cpy(ctx, src0, node); + return ggml_webgpu_cpy(ctx, encoder, src0, node); case GGML_OP_SET: - return ggml_webgpu_set(ctx, src0, src1, node); + return ggml_webgpu_set(ctx, encoder, src0, src1, node); case GGML_OP_SET_ROWS: - return ggml_webgpu_set_rows(ctx, src0, src1, node); + return ggml_webgpu_set_rows(ctx, encoder, src0, src1, node); case GGML_OP_GET_ROWS: - return ggml_webgpu_get_rows(ctx, src0, src1, node); + return ggml_webgpu_get_rows(ctx, encoder, src0, src1, node); case GGML_OP_MUL_MAT: - return ggml_webgpu_mul_mat(ctx, src0, src1, node); + return ggml_webgpu_mul_mat(ctx, encoder, src0, src1, node); case GGML_OP_FLASH_ATTN_EXT: #ifndef __EMSCRIPTEN__ - return ggml_webgpu_flash_attn(ctx, src0, src1, src2, node->src[3], node->src[4], node); + return ggml_webgpu_flash_attn(ctx, encoder, src0, src1, src2, node->src[3], node->src[4], node); #else return std::nullopt; #endif @@ -2391,22 +2387,22 @@ static std::optional ggml_webgpu_encode_node(webgpu_context ctx, case GGML_OP_SUB: case GGML_OP_MUL: case GGML_OP_DIV: - return ggml_webgpu_binary_op(ctx, src0, src1, node); + return ggml_webgpu_binary_op(ctx, encoder, src0, src1, node); case GGML_OP_CONCAT: - return ggml_webgpu_concat(ctx, src0, src1, node); + return ggml_webgpu_concat(ctx, encoder, src0, src1, node); case GGML_OP_REPEAT: - return ggml_webgpu_repeat(ctx, src0, node); + return ggml_webgpu_repeat(ctx, encoder, src0, node); case GGML_OP_RMS_NORM: case GGML_OP_L2_NORM: - return ggml_webgpu_row_norm(ctx, src0, node); + return ggml_webgpu_row_norm(ctx, encoder, src0, node); case GGML_OP_ROPE: - return ggml_webgpu_rope(ctx, src0, src1, src2, node); + return ggml_webgpu_rope(ctx, encoder, src0, src1, src2, node); case GGML_OP_GLU: - return ggml_webgpu_glu(ctx, src0, src1, node); + return ggml_webgpu_glu(ctx, encoder, src0, src1, node); case GGML_OP_SCALE: - return ggml_webgpu_scale(ctx, src0, node); + return ggml_webgpu_scale(ctx, encoder, src0, node); case GGML_OP_SOFT_MAX: - return ggml_webgpu_soft_max(ctx, src0, src1, src2, node); + return ggml_webgpu_soft_max(ctx, encoder, src0, src1, src2, node); case GGML_OP_UNARY: case GGML_OP_CLAMP: case GGML_OP_FILL: @@ -2417,26 +2413,26 @@ static std::optional ggml_webgpu_encode_node(webgpu_context ctx, case GGML_OP_COS: case GGML_OP_DIAG: case GGML_OP_TRI: - return ggml_webgpu_unary_op(ctx, src0, node); + return ggml_webgpu_unary_op(ctx, encoder, src0, node); case GGML_OP_SOLVE_TRI: - return ggml_webgpu_solve_tri(ctx, src0, src1, node); + return ggml_webgpu_solve_tri(ctx, encoder, src0, src1, node); case GGML_OP_SSM_CONV: - return ggml_webgpu_ssm_conv(ctx, src0, src1, node); + return ggml_webgpu_ssm_conv(ctx, encoder, src0, src1, node); case GGML_OP_GATED_DELTA_NET: - return ggml_webgpu_gated_delta_net(ctx, src0, src1, src2, node->src[3], node->src[4], node->src[5], node); + return ggml_webgpu_gated_delta_net(ctx, encoder, src0, src1, src2, node->src[3], node->src[4], node->src[5], node); case GGML_OP_PAD: - return ggml_webgpu_pad(ctx, src0, node); + return ggml_webgpu_pad(ctx, encoder, src0, node); case GGML_OP_ARGMAX: - return ggml_webgpu_argmax(ctx, src0, node); + return ggml_webgpu_argmax(ctx, encoder, src0, node); case GGML_OP_ARGSORT: case GGML_OP_TOP_K: // we reuse the same argsort implementation for top_k - return ggml_webgpu_argsort(ctx, src0, node); + return ggml_webgpu_argsort(ctx, encoder, src0, node); case GGML_OP_CUMSUM: - return ggml_webgpu_cumsum(ctx, src0, node); + return ggml_webgpu_cumsum(ctx, encoder, src0, node); case GGML_OP_SUM: case GGML_OP_SUM_ROWS: - return ggml_webgpu_sum_rows(ctx, src0, node); + return ggml_webgpu_sum_rows(ctx, encoder, src0, node); default: return std::nullopt; } @@ -2450,34 +2446,38 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str WEBGPU_CPU_PROFILE_TOTAL_START(graph_compute); - std::vector commands; + std::vector commands; #ifdef GGML_WEBGPU_GPU_PROFILE std::vector profile_futures; #endif uint32_t num_batched_kernels = 0; bool contains_set_rows = false; + wgpu::CommandEncoder batch_encoder = ctx->global_ctx->device.CreateCommandEncoder(); for (int i = 0; i < cgraph->n_nodes; i++) { if (cgraph->nodes[i]->op == GGML_OP_SET_ROWS) { contains_set_rows = true; } - if (auto cmd = ggml_webgpu_encode_node(ctx, cgraph->nodes[i])) { + if (auto cmd = ggml_webgpu_encode_node(ctx, batch_encoder, cgraph->nodes[i])) { commands.push_back(*cmd); num_batched_kernels += cmd.value().num_kernels; } if (num_batched_kernels >= WEBGPU_COMMAND_SUBMIT_BATCH_SIZE) { num_batched_kernels = 0; - ggml_backend_webgpu_submit(ctx->global_ctx, commands); + wgpu::CommandBuffer batch_commands = batch_encoder.Finish(); + ctx->global_ctx->queue.Submit(1, &batch_commands); #ifdef GGML_WEBGPU_GPU_PROFILE ggml_backend_webgpu_collect_profile_futures(ctx->global_ctx, commands, profile_futures); #endif ctx->param_arena.reset(); commands.clear(); + batch_encoder = ctx->global_ctx->device.CreateCommandEncoder(); } } if (!commands.empty()) { - ggml_backend_webgpu_submit(ctx->global_ctx, commands); + wgpu::CommandBuffer batch_commands = batch_encoder.Finish(); + ctx->global_ctx->queue.Submit(1, &batch_commands); #ifdef GGML_WEBGPU_GPU_PROFILE ggml_backend_webgpu_collect_profile_futures(ctx->global_ctx, commands, profile_futures); #endif From 7450dd01c0268f8117e08a54c8e0cee91b462040 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Fri, 3 Apr 2026 08:08:14 -0700 Subject: [PATCH 12/16] formatting --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 43 +++++++++++-------- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 4 +- 2 files changed, 26 insertions(+), 21 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index 1c56c689312f..669d2cd53a83 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -437,12 +437,18 @@ inline uint32_t ggml_webgpu_flash_attn_pick_vec_ne(const ggml_webgpu_flash_attn_ // Head-dim specializations used by the tuned vec f16 path. switch (key.head_dim_qk) { - case 64: return 2u; - case 96: return 4u; - case 128: return 1u; - case 192: return 2u; - case 576: return 2u; - default: return 1u; + case 64: + return 2u; + case 96: + return 4u; + case 128: + return 1u; + case 192: + return 2u; + case 576: + return 2u; + default: + return 1u; } } @@ -513,9 +519,9 @@ struct ggml_webgpu_flash_attn_blk_shader_lib_context { }; inline ggml_webgpu_processed_shader ggml_webgpu_preprocess_flash_attn_blk_shader( - pre_wgsl::Preprocessor & preprocessor, - const char * shader_src, - const ggml_webgpu_flash_attn_blk_shader_lib_context & context) { + pre_wgsl::Preprocessor & preprocessor, + const char * shader_src, + const ggml_webgpu_flash_attn_blk_shader_lib_context & context) { std::vector defines; std::string variant = "flash_attn_vec_blk"; @@ -1857,9 +1863,8 @@ class ggml_webgpu_shader_lib { defines.push_back(std::string("SG_MAT_K=") + std::to_string(context.sg_mat_k)); uint32_t q_tile = context.sg_mat_m; - uint32_t kv_tile = - std::min(ggml_webgpu_flash_attn_max_kv_tile(context), - context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES); + uint32_t kv_tile = std::min(ggml_webgpu_flash_attn_max_kv_tile(context), + context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES); if (context.key.use_vec) { q_tile = 1; kv_tile = std::max(context.sg_mat_n, std::min(32u, ggml_webgpu_flash_attn_max_kv_tile(context))); @@ -1885,14 +1890,14 @@ class ggml_webgpu_shader_lib { } defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - const char * shader_src = context.key.use_vec ? wgsl_flash_attn_vec_split : wgsl_flash_attn; + const char * shader_src = context.key.use_vec ? wgsl_flash_attn_vec_split : wgsl_flash_attn; webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, preprocessor.preprocess(shader_src, defines), variant); - auto decisions = std::make_shared(); - decisions->q_tile = q_tile; - decisions->kv_tile = kv_tile; - decisions->wg_size = wg_size; - pipeline.context = decisions; + auto decisions = std::make_shared(); + decisions->q_tile = q_tile; + decisions->kv_tile = kv_tile; + decisions->wg_size = wg_size; + pipeline.context = decisions; flash_attn_pipelines[context.key] = pipeline; return flash_attn_pipelines[context.key]; } @@ -1905,7 +1910,7 @@ class ggml_webgpu_shader_lib { ggml_webgpu_processed_shader processed = ggml_webgpu_preprocess_flash_attn_blk_shader(preprocessor, wgsl_flash_attn_vec_blk, context); - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed.wgsl, processed.variant); + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed.wgsl, processed.variant); flash_attn_blk_pipelines[context.key] = pipeline; return flash_attn_blk_pipelines[context.key]; } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index ac87195f6296..5c567dc0df07 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -2709,9 +2709,9 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str WEBGPU_CPU_PROFILE_TOTAL_START(graph_compute); std::vector commands; -# ifdef GGML_WEBGPU_GPU_PROFILE +#ifdef GGML_WEBGPU_GPU_PROFILE std::vector profile_futures; -# endif +#endif uint32_t num_batched_kernels = 0; bool contains_set_rows = false; wgpu::CommandEncoder batch_encoder = ctx->global_ctx->device.CreateCommandEncoder(); From cc6f8089f0fbbf02796a87f3c608e55c1b22b12e Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Fri, 3 Apr 2026 11:09:17 -0700 Subject: [PATCH 13/16] Add parameters for different browsers in-flight submissions --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 118 ++++++++++++++++++++------- 1 file changed, 90 insertions(+), 28 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 5c567dc0df07..1853796d095d 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -11,12 +11,17 @@ #ifdef __EMSCRIPTEN__ # include + +// iOS browsers have very strict limits on the number of in-flight GPU commands, so we need to limit to 1 batch in flight at a time to avoid GPU timeouts. +EM_JS(int, ggml_webgpu_is_ios_browser, (), { + const ua = navigator.userAgent; + return (ua.includes('iPhone') || ua.includes('iPad')) ? 1 : 0; +}); #endif #include #include -#include #include #include #ifdef GGML_WEBGPU_GPU_PROFILE @@ -25,7 +30,6 @@ #if defined(GGML_WEBGPU_DEBUG) || defined(GGML_WEBGPU_CPU_PROFILE) || defined(GGML_WEBGPU_GPU_PROFILE) # include #endif -#include #include #include #include @@ -84,7 +88,8 @@ static inline void compute_2d_workgroups(uint32_t total_wg, uint32_t max_per_dim #define WEBGPU_COMMAND_SUBMIT_BATCH_SIZE 32u #define WEBGPU_NUM_PARAM_SLOTS \ (WEBGPU_COMMAND_SUBMIT_BATCH_SIZE + 10) // a few extra for safety, since some operations may need multiple slots -#define WEBGPU_WAIT_ANY_TIMEOUT_MS 100 +#define WEBGPU_RUNTIME_WAIT_TIMEOUT_MS 5000u +#define WEBGPU_RUNTIME_WAIT_TIMEOUT_NS (WEBGPU_RUNTIME_WAIT_TIMEOUT_MS * 1e6) #define WEBGPU_PARAMS_BUF_SIZE_BYTES 128 // enough for 32 parameters #define WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES 4 #define WEBGPU_STORAGE_BUF_BINDING_MULT 4 // a storage buffer binding size must be a multiple of 4 @@ -252,6 +257,7 @@ struct webgpu_global_context_struct { wgpu::Adapter adapter; wgpu::Device device; wgpu::Queue queue; + uint32_t max_inflight_batches = UINT32_MAX; webgpu_capabilities capabilities; // Shared buffer to move data from device to host @@ -417,16 +423,52 @@ static void ggml_backend_webgpu_wait_profile_futures(webgpu_global_context & } #endif +template +static void ggml_backend_webgpu_check_wait_status(wgpu::WaitStatus wait_status, + T callback_status, + T success_status, + const char * wait_name, + const char * failure_name, + const char * callback_message) { + if (wait_status == wgpu::WaitStatus::TimedOut) { + GGML_ABORT("ggml_webgpu: %s timed out after %u ms\n", wait_name, WEBGPU_RUNTIME_WAIT_TIMEOUT_MS); + } + if (wait_status == wgpu::WaitStatus::Error) { + GGML_ABORT("ggml_webgpu: %s failed\n", wait_name); + } + if (callback_status != success_status) { + GGML_ABORT("ggml_webgpu: %s failed with status %d: %s\n", failure_name, static_cast(callback_status), + callback_message); + } +} + +static uint32_t ggml_backend_webgpu_get_max_inflight_batches(const wgpu::AdapterInfo & info) { +#ifdef __EMSCRIPTEN__ + if (ggml_webgpu_is_ios_browser()) { + return 1; + } +#else + GGML_UNUSED(info); +#endif + + return UINT32_MAX; +} + static void ggml_backend_webgpu_wait_queue(webgpu_global_context & ctx) { - ctx->instance.WaitAny( - ctx->queue.OnSubmittedWorkDone(wgpu::CallbackMode::AllowSpontaneous, - [](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { - if (status != wgpu::QueueWorkDoneStatus::Success) { - GGML_LOG_ERROR("ggml_webgpu: Failed to submit commands: %s\n", - std::string(message).c_str()); - } - }), - UINT64_MAX); + wgpu::QueueWorkDoneStatus callback_status = wgpu::QueueWorkDoneStatus::Error; + std::string callback_message; + + const wgpu::WaitStatus wait_status = ctx->instance.WaitAny( + ctx->queue.OnSubmittedWorkDone( + wgpu::CallbackMode::AllowSpontaneous, + [&callback_status, &callback_message](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { + callback_status = status; + callback_message = std::string(message); + }), + WEBGPU_RUNTIME_WAIT_TIMEOUT_NS); + + ggml_backend_webgpu_check_wait_status(wait_status, callback_status, wgpu::QueueWorkDoneStatus::Success, + "Queue wait", "Queue work", callback_message.c_str()); } static void ggml_backend_webgpu_map_buffer(webgpu_global_context & ctx, @@ -434,14 +476,31 @@ static void ggml_backend_webgpu_map_buffer(webgpu_global_context & ctx, wgpu::MapMode mode, size_t offset, size_t size) { - ctx->instance.WaitAny(buffer.MapAsync(mode, offset, size, wgpu::CallbackMode::AllowSpontaneous, - [](wgpu::MapAsyncStatus status, wgpu::StringView message) { - if (status != wgpu::MapAsyncStatus::Success) { - GGML_LOG_ERROR("ggml_webgpu: Failed to map buffer: %s\n", - message.data); - } - }), - UINT64_MAX); + wgpu::MapAsyncStatus callback_status = wgpu::MapAsyncStatus::Error; + std::string callback_message; + + const wgpu::WaitStatus wait_status = ctx->instance.WaitAny( + buffer.MapAsync(mode, offset, size, wgpu::CallbackMode::AllowSpontaneous, + [&callback_status, &callback_message](wgpu::MapAsyncStatus status, wgpu::StringView message) { + callback_status = status; + callback_message = std::string(message); + }), + WEBGPU_RUNTIME_WAIT_TIMEOUT_NS); + + ggml_backend_webgpu_check_wait_status(wait_status, callback_status, wgpu::MapAsyncStatus::Success, + "Buffer map wait", "Buffer map", callback_message.c_str()); +} + +static void ggml_backend_webgpu_submit_commands(webgpu_context & ctx, + const wgpu::CommandBuffer commands, + uint32_t & num_inflight_batches) { + if (num_inflight_batches >= ctx->global_ctx->max_inflight_batches) { + ggml_backend_webgpu_wait_queue(ctx->global_ctx); + num_inflight_batches = 0; + } + + ctx->global_ctx->queue.Submit(1, &commands); + num_inflight_batches++; } #ifdef GGML_WEBGPU_DEBUG @@ -2712,9 +2771,10 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str #ifdef GGML_WEBGPU_GPU_PROFILE std::vector profile_futures; #endif - uint32_t num_batched_kernels = 0; - bool contains_set_rows = false; - wgpu::CommandEncoder batch_encoder = ctx->global_ctx->device.CreateCommandEncoder(); + uint32_t num_batched_kernels = 0; + uint32_t num_inflight_batches = 0; + bool contains_set_rows = false; + wgpu::CommandEncoder batch_encoder = ctx->global_ctx->device.CreateCommandEncoder(); for (int i = 0; i < cgraph->n_nodes; i++) { if (cgraph->nodes[i]->op == GGML_OP_SET_ROWS) { @@ -2728,7 +2788,7 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str if (num_batched_kernels >= WEBGPU_COMMAND_SUBMIT_BATCH_SIZE) { num_batched_kernels = 0; wgpu::CommandBuffer batch_commands = batch_encoder.Finish(); - ctx->global_ctx->queue.Submit(1, &batch_commands); + ggml_backend_webgpu_submit_commands(ctx, batch_commands, num_inflight_batches); #ifdef GGML_WEBGPU_GPU_PROFILE ggml_backend_webgpu_collect_profile_futures(ctx->global_ctx, commands, profile_futures); #endif @@ -2739,7 +2799,7 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str } if (!commands.empty()) { wgpu::CommandBuffer batch_commands = batch_encoder.Finish(); - ctx->global_ctx->queue.Submit(1, &batch_commands); + ggml_backend_webgpu_submit_commands(ctx, batch_commands, num_inflight_batches); #ifdef GGML_WEBGPU_GPU_PROFILE ggml_backend_webgpu_collect_profile_futures(ctx->global_ctx, commands, profile_futures); #endif @@ -2753,7 +2813,7 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str encoder.CopyBufferToBuffer(ctx->set_rows_dev_error_buf, 0, ctx->set_rows_host_error_buf, 0, ctx->set_rows_host_error_buf.GetSize()); wgpu::CommandBuffer set_rows_commands = encoder.Finish(); - ctx->global_ctx->queue.Submit(1, &set_rows_commands); + ggml_backend_webgpu_submit_commands(ctx, set_rows_commands, num_inflight_batches); } ggml_backend_webgpu_wait_queue(ctx->global_ctx); @@ -3190,6 +3250,7 @@ static bool create_webgpu_device(ggml_backend_webgpu_reg_context * ctx) { } #endif ctx->webgpu_global_ctx->adapter.GetInfo(&info); + ctx->webgpu_global_ctx->max_inflight_batches = ggml_backend_webgpu_get_max_inflight_batches(info); wgpu::SupportedFeatures features; ctx->webgpu_global_ctx->adapter.GetFeatures(&features); // we require f16 support @@ -3299,9 +3360,10 @@ static bool create_webgpu_device(ggml_backend_webgpu_reg_context * ctx) { GGML_LOG_INFO( "ggml_webgpu: adapter_info: vendor_id: %u | vendor: %s | architecture: %s | device_id: %u | name: %s | " - "device_desc: %s\n", + "device_desc: %s | max_inflight_batches: %u\n", info.vendorID, std::string(info.vendor).c_str(), std::string(info.architecture).c_str(), info.deviceID, - std::string(info.device).c_str(), std::string(info.description).c_str()); + std::string(info.device).c_str(), std::string(info.description).c_str(), + ctx->webgpu_global_ctx->max_inflight_batches); return true; } From d63c7395e7a4bc8002b11650f2bce5d0491fedf4 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Fri, 3 Apr 2026 11:33:02 -0700 Subject: [PATCH 14/16] Update handling of batch size too --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 62 +++++++++++++++++----------- 1 file changed, 39 insertions(+), 23 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 1853796d095d..3898c66b77cf 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -11,12 +11,6 @@ #ifdef __EMSCRIPTEN__ # include - -// iOS browsers have very strict limits on the number of in-flight GPU commands, so we need to limit to 1 batch in flight at a time to avoid GPU timeouts. -EM_JS(int, ggml_webgpu_is_ios_browser, (), { - const ua = navigator.userAgent; - return (ua.includes('iPhone') || ua.includes('iPad')) ? 1 : 0; -}); #endif #include @@ -85,14 +79,13 @@ static inline void compute_2d_workgroups(uint32_t total_wg, uint32_t max_per_dim /* Constants */ -#define WEBGPU_COMMAND_SUBMIT_BATCH_SIZE 32u -#define WEBGPU_NUM_PARAM_SLOTS \ - (WEBGPU_COMMAND_SUBMIT_BATCH_SIZE + 10) // a few extra for safety, since some operations may need multiple slots -#define WEBGPU_RUNTIME_WAIT_TIMEOUT_MS 5000u -#define WEBGPU_RUNTIME_WAIT_TIMEOUT_NS (WEBGPU_RUNTIME_WAIT_TIMEOUT_MS * 1e6) -#define WEBGPU_PARAMS_BUF_SIZE_BYTES 128 // enough for 32 parameters -#define WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES 4 -#define WEBGPU_STORAGE_BUF_BINDING_MULT 4 // a storage buffer binding size must be a multiple of 4 +#define WEBGPU_DEFAULT_COMMAND_SUBMIT_BATCH_SIZE 32u +#define WEBGPU_NUM_PARAM_SLOT_SAFETY_MARGIN 10u +#define WEBGPU_RUNTIME_WAIT_TIMEOUT_MS 5000u +#define WEBGPU_RUNTIME_WAIT_TIMEOUT_NS (WEBGPU_RUNTIME_WAIT_TIMEOUT_MS * 1e6) +#define WEBGPU_PARAMS_BUF_SIZE_BYTES 128 // enough for 32 parameters +#define WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES 4 +#define WEBGPU_STORAGE_BUF_BINDING_MULT 4 // a storage buffer binding size must be a multiple of 4 // For operations which process a row in parallel, this seems like a reasonable // default @@ -257,7 +250,8 @@ struct webgpu_global_context_struct { wgpu::Adapter adapter; wgpu::Device device; wgpu::Queue queue; - uint32_t max_inflight_batches = UINT32_MAX; + uint32_t command_submit_batch_size = WEBGPU_DEFAULT_COMMAND_SUBMIT_BATCH_SIZE; + uint32_t max_inflight_batches = UINT32_MAX; webgpu_capabilities capabilities; // Shared buffer to move data from device to host @@ -442,10 +436,18 @@ static void ggml_backend_webgpu_check_wait_status(wgpu::WaitStatus wait_status, } } +#ifdef __EMSCRIPTEN__ +// iOS browsers seem to have very strict limits on the number of in-flight GPU commands, so we need to throttle to avoid failures. +EM_JS(int, ggml_webgpu_is_ios_browser, (), { + const ua = navigator.userAgent; + return (ua.includes('iPhone') || ua.includes('iPad')) ? 1 : 0; +}); +#endif + static uint32_t ggml_backend_webgpu_get_max_inflight_batches(const wgpu::AdapterInfo & info) { #ifdef __EMSCRIPTEN__ if (ggml_webgpu_is_ios_browser()) { - return 1; + return 2; } #else GGML_UNUSED(info); @@ -454,6 +456,18 @@ static uint32_t ggml_backend_webgpu_get_max_inflight_batches(const wgpu::Adapter return UINT32_MAX; } +static uint32_t ggml_backend_webgpu_get_command_submit_batch_size(const wgpu::AdapterInfo & info) { +#ifdef __EMSCRIPTEN__ + if (ggml_webgpu_is_ios_browser()) { + return 16; + } +#else + GGML_UNUSED(info); +#endif + + return WEBGPU_DEFAULT_COMMAND_SUBMIT_BATCH_SIZE; +} + static void ggml_backend_webgpu_wait_queue(webgpu_global_context & ctx) { wgpu::QueueWorkDoneStatus callback_status = wgpu::QueueWorkDoneStatus::Error; std::string callback_message; @@ -2785,7 +2799,7 @@ static ggml_status ggml_backend_webgpu_graph_compute(ggml_backend_t backend, str num_batched_kernels += cmd.value().num_kernels; } - if (num_batched_kernels >= WEBGPU_COMMAND_SUBMIT_BATCH_SIZE) { + if (num_batched_kernels >= ctx->global_ctx->command_submit_batch_size) { num_batched_kernels = 0; wgpu::CommandBuffer batch_commands = batch_encoder.Finish(); ggml_backend_webgpu_submit_commands(ctx, batch_commands, num_inflight_batches); @@ -3250,7 +3264,8 @@ static bool create_webgpu_device(ggml_backend_webgpu_reg_context * ctx) { } #endif ctx->webgpu_global_ctx->adapter.GetInfo(&info); - ctx->webgpu_global_ctx->max_inflight_batches = ggml_backend_webgpu_get_max_inflight_batches(info); + ctx->webgpu_global_ctx->command_submit_batch_size = ggml_backend_webgpu_get_command_submit_batch_size(info); + ctx->webgpu_global_ctx->max_inflight_batches = ggml_backend_webgpu_get_max_inflight_batches(info); wgpu::SupportedFeatures features; ctx->webgpu_global_ctx->adapter.GetFeatures(&features); // we require f16 support @@ -3360,10 +3375,9 @@ static bool create_webgpu_device(ggml_backend_webgpu_reg_context * ctx) { GGML_LOG_INFO( "ggml_webgpu: adapter_info: vendor_id: %u | vendor: %s | architecture: %s | device_id: %u | name: %s | " - "device_desc: %s | max_inflight_batches: %u\n", + "device_desc: %s\n", info.vendorID, std::string(info.vendor).c_str(), std::string(info.architecture).c_str(), info.deviceID, - std::string(info.device).c_str(), std::string(info.description).c_str(), - ctx->webgpu_global_ctx->max_inflight_batches); + std::string(info.device).c_str(), std::string(info.description).c_str()); return true; } @@ -3372,8 +3386,10 @@ static webgpu_context initialize_webgpu_context(ggml_backend_dev_t dev) { webgpu_context webgpu_ctx = std::make_shared(); webgpu_ctx->global_ctx = dev_ctx->webgpu_global_ctx; webgpu_ctx->shader_lib = std::make_unique(dev_ctx->webgpu_global_ctx->device); - webgpu_ctx->param_arena.init(webgpu_ctx->global_ctx->device, WEBGPU_PARAMS_BUF_SIZE_BYTES, WEBGPU_NUM_PARAM_SLOTS, - webgpu_ctx->global_ctx->capabilities.limits.minUniformBufferOffsetAlignment); + webgpu_ctx->param_arena.init( + webgpu_ctx->global_ctx->device, WEBGPU_PARAMS_BUF_SIZE_BYTES, + webgpu_ctx->global_ctx->command_submit_batch_size + WEBGPU_NUM_PARAM_SLOT_SAFETY_MARGIN, + webgpu_ctx->global_ctx->capabilities.limits.minUniformBufferOffsetAlignment); ggml_webgpu_create_buffer(webgpu_ctx->global_ctx->device, webgpu_ctx->set_rows_dev_error_buf, WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES, wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc, "set_rows_dev_error_buf"); From f9d63dc5673fab4fdcba03dede66f01571706c72 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Fri, 3 Apr 2026 12:31:48 -0700 Subject: [PATCH 15/16] Throttle ios as much as possible --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 3898c66b77cf..f391c3bb8d84 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -447,7 +447,7 @@ EM_JS(int, ggml_webgpu_is_ios_browser, (), { static uint32_t ggml_backend_webgpu_get_max_inflight_batches(const wgpu::AdapterInfo & info) { #ifdef __EMSCRIPTEN__ if (ggml_webgpu_is_ios_browser()) { - return 2; + return 1; } #else GGML_UNUSED(info); From 5b0fd3d29bbfd439800f4bed4599e82228016c47 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Mon, 6 Apr 2026 16:25:16 -0700 Subject: [PATCH 16/16] Increase timeout for llvm-pipe testing --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index ea568c90ccee..3d038924b788 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -81,7 +81,7 @@ static inline void compute_2d_workgroups(uint32_t total_wg, uint32_t max_per_dim #define WEBGPU_DEFAULT_COMMAND_SUBMIT_BATCH_SIZE 32u #define WEBGPU_NUM_PARAM_SLOT_SAFETY_MARGIN 10u -#define WEBGPU_RUNTIME_WAIT_TIMEOUT_MS 5000u +#define WEBGPU_RUNTIME_WAIT_TIMEOUT_MS 30000u #define WEBGPU_RUNTIME_WAIT_TIMEOUT_NS (WEBGPU_RUNTIME_WAIT_TIMEOUT_MS * 1e6) #define WEBGPU_PARAMS_BUF_SIZE_BYTES 128 // enough for 32 parameters #define WEBGPU_SET_ROWS_ERROR_BUF_SIZE_BYTES 4