ggml webgpu: ops support for qwen3.5 (SET, TRI_SOLVE, SSM_CONV, GATED_DELTA_NET) + GET_ROWS optimization (#20687)

* Implement l2_norm, set, tri

* Add DIAG/SOLVE_TRI

* Add SSM_CONV

* Better get_rows and gated_delta_net to support qwen3.5

* Clean up, update ops.md

* Fix binding_index type for wasm

* Fix read write annotations

* cleanups

ggml/src/ggml-webgpu/wgsl-shaders/gated_delta_net.wgsl

@@ -0,0 +1,132 @@
+@group(0) @binding(0)
+var<storage, read_write> src_q: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> src_k: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read_write> src_v: array<f32>;
+
+@group(0) @binding(3)
+var<storage, read_write> src_g: array<f32>;
+
+@group(0) @binding(4)
+var<storage, read_write> src_beta: array<f32>;
+
+@group(0) @binding(5)
+var<storage, read_write> src_state: array<f32>;
+
+@group(0) @binding(6)
+var<storage, read_write> dst: array<f32>;
+
+struct Params {
+    h: u32,
+    n_tokens: u32,
+    n_seqs: u32,
+    s_off: u32,
+
+    sq1: u32,
+    sq2: u32,
+    sq3: u32,
+
+    sv1: u32,
+    sv2: u32,
+    sv3: u32,
+
+    sb1: u32,
+    sb2: u32,
+    sb3: u32,
+
+    neq1: u32,
+    rq3: u32,
+    scale: f32,
+};
+
+@group(0) @binding(7)
+var<uniform> params: Params;
+
+var<workgroup> sh_k: array<f32, S_V>;
+var<workgroup> sh_q: array<f32, S_V>;
+var<workgroup> sh_g: array<f32, S_V>;
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+    @builtin(workgroup_id) workgroup_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>
+) {
+    let head_id = workgroup_id.x;
+    let seq_id = workgroup_id.y;
+    let col = local_id.x;
+
+    let iq1 = head_id % params.neq1;
+    let iq3 = seq_id / params.rq3;
+
+    let state_size = S_V * S_V;
+    let state_base = (seq_id * params.h + head_id) * state_size;
+
+    var state: array<f32, S_V>;
+    for (var i = 0u; i < S_V; i++) {
+        state[i] = src_state[state_base + col * S_V + i];
+    }
+
+    var attn_off = (seq_id * params.n_tokens * params.h + head_id) * S_V;
+
+    for (var t = 0u; t < params.n_tokens; t++) {
+        let q_off = iq3 * params.sq3 + t * params.sq2 + iq1 * params.sq1;
+        let k_off = q_off;
+        let v_off = seq_id * params.sv3 + t * params.sv2 + head_id * params.sv1;
+        let gb_off = seq_id * params.sb3 + t * params.sb2 + head_id * params.sb1;
+
+        sh_q[col] = src_q[q_off + col];
+        sh_k[col] = src_k[k_off + col];
+
+#ifdef KDA
+        let g_base = gb_off * S_V;
+        sh_g[col] = exp(src_g[g_base + col]);
+#endif
+
+        workgroupBarrier();
+
+        let v_val = src_v[v_off + col];
+        let beta_val = src_beta[gb_off];
+
+        var kv_col = 0.0;
+        var delta_col = 0.0;
+        var attn_col = 0.0;
+
+#ifdef KDA
+        for (var i = 0u; i < S_V; i++) {
+            kv_col += (sh_g[i] * state[i]) * sh_k[i];
+        }
+
+        delta_col = (v_val - kv_col) * beta_val;
+
+        for (var i = 0u; i < S_V; i++) {
+            state[i] = sh_g[i] * state[i] + sh_k[i] * delta_col;
+            attn_col += state[i] * sh_q[i];
+        }
+#else
+        let g_val = exp(src_g[gb_off]);
+
+        for (var i = 0u; i < S_V; i++) {
+            kv_col += state[i] * sh_k[i];
+        }
+
+        delta_col = (v_val - g_val * kv_col) * beta_val;
+
+        for (var i = 0u; i < S_V; i++) {
+            state[i] = g_val * state[i] + sh_k[i] * delta_col;
+            attn_col += state[i] * sh_q[i];
+        }
+#endif
+
+        dst[attn_off + col] = attn_col * params.scale;
+        attn_off += S_V * params.h;
+
+        workgroupBarrier();
+    }
+
+    for (var i = 0u; i < S_V; i++) {
+        dst[params.s_off + state_base + col * S_V + i] = state[i];
+    }
+}

ggml/src/ggml-webgpu/wgsl-shaders/get_rows.wgsl

@@ -640,6 +640,35 @@ var<uniform> params: Params;
 
 @compute @workgroup_size(WG_SIZE)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+#ifdef FLOAT_PARALLEL
+    let blocks_per_row = params.ne0 / BLOCK_SIZE;
+    let row_count = params.n_rows * params.ne2 * params.ne3;
+
+    if (gid.x >= blocks_per_row * row_count) {
+        return;
+    }
+
+    let block_idx = gid.x % blocks_per_row;
+    var row_idx = gid.x / blocks_per_row;
+    let i_dst3 = row_idx / (params.ne2 * params.n_rows);
+
+    row_idx = row_idx % (params.ne2 * params.n_rows);
+    let i_dst2 = row_idx / params.n_rows;
+    let i_dst1 = row_idx % params.n_rows;
+
+    let i_idx2 = i_dst3 % params.idx2;
+    let i_idx1 = i_dst2 % params.idx1;
+    let i_idx0 = i_dst1;
+
+    let i_idx = params.offset_idx + i_idx0 * params.stride_idx0 + i_idx1 * params.stride_idx1 + i_idx2 * params.stride_idx2;
+
+    let idx_val = u32(idx[i_idx]);
+
+    let i_src_row = params.offset_src + idx_val * params.stride_src1 + i_dst2 * params.stride_src2 + i_dst3 * params.stride_src3;
+    let i_dst_row = params.offset_dst + i_dst1 * params.stride_dst1 + i_dst2 * params.stride_dst2 + i_dst3 * params.stride_dst3;
+
+    copy_elements(i_src_row, i_dst_row, block_idx);
+#else
     if (gid.x >= params.n_rows * params.ne2 * params.ne3) {
         return;
     }
@@ -664,5 +693,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     for (var i: u32 = 0; i < params.ne0/BLOCK_SIZE; i++) {
       copy_elements(i_src_row, i_dst_row, i);
     }
+#endif
 }
-

ggml/src/ggml-webgpu/wgsl-shaders/row_norm.wgsl

@@ -81,11 +81,12 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
     }
     sum = scratch[0];
 
-#ifdef OP_RMS_NORM
+#ifdef RMS_NORM
     let scale = 1.0/sqrt(sum/f32(params.ne0) + params.eps);
-#elif OP_L2_NORM
+#elif defined(L2_NORM)
     let scale = 1.0/max(sqrt(sum), params.eps);
 #endif
+
     col = lid.x;
     for (var j: u32 = 0; j < elems; j++) {
         if (col >= params.ne0) {

ggml/src/ggml-webgpu/wgsl-shaders/set.wgsl

@@ -0,0 +1,109 @@
+#ifdef TYPE_I32
+#define TYPE i32
+#else
+#define TYPE f32
+#endif
+
+#ifndef INPLACE
+@group(0) @binding(0)
+var<storage, read_write> src0: array<TYPE>;
+#define SRC1_BINDING 1
+#else
+#define SRC1_BINDING 0
+#endif
+
+#define DST_BINDING SRC1_BINDING + 1
+#define PARAMS_BINDING SRC1_BINDING + 2
+
+@group(0) @binding(SRC1_BINDING)
+var<storage, read_write> src1: array<TYPE>;
+
+@group(0) @binding(DST_BINDING)
+var<storage, read_write> dst: array<TYPE>;
+
+struct Params {
+    ne: u32,
+    offset_src0: u32,
+    offset_src1: u32,
+    offset_view: u32,
+
+    stride_src10: u32,
+    stride_src11: u32,
+    stride_src12: u32,
+    stride_src13: u32,
+
+    stride_dst10: u32,
+    stride_dst11: u32,
+    stride_dst12: u32,
+    stride_dst13: u32,
+
+    src1_ne0: u32,
+    src1_ne1: u32,
+    src1_ne2: u32,
+    src1_ne3: u32,
+};
+
+@group(0) @binding(PARAMS_BINDING)
+var<uniform> params: Params;
+
+fn decode_src1_coords(idx: u32) -> vec4<u32> {
+    var i = idx;
+    let plane = params.src1_ne2 * params.src1_ne1 * params.src1_ne0;
+    let i3 = i / plane;
+    i = i % plane;
+    let row = params.src1_ne1 * params.src1_ne0;
+    let i2 = i / row;
+    i = i % row;
+    let i1 = i / params.src1_ne0;
+    let i0 = i % params.src1_ne0;
+    return vec4<u32>(i0, i1, i2, i3);
+}
+
+fn decode_view_coords(rel: u32) -> vec4<u32> {
+    let i3 = rel / params.stride_dst13;
+    let rem3 = rel % params.stride_dst13;
+    let i2 = rem3 / params.stride_dst12;
+    let rem2 = rem3 % params.stride_dst12;
+    let i1 = rem2 / params.stride_dst11;
+    let i0 = rem2 % params.stride_dst11;
+    return vec4<u32>(i0, i1, i2, i3);
+}
+
+fn view_rel_from_coords(coords: vec4<u32>) -> u32 {
+    return coords.x * params.stride_dst10 + coords.y * params.stride_dst11 +
+           coords.z * params.stride_dst12 + coords.w * params.stride_dst13;
+}
+
+fn src1_idx_from_coords(coords: vec4<u32>) -> u32 {
+    return coords.x * params.stride_src10 + coords.y * params.stride_src11 +
+           coords.z * params.stride_src12 + coords.w * params.stride_src13;
+}
+
+fn in_set_view(rel: u32, coords: vec4<u32>) -> bool {
+    return view_rel_from_coords(coords) == rel;
+}
+
+@compute @workgroup_size(WG_SIZE)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    if (gid.x >= params.ne) {
+        return;
+    }
+
+#ifdef INPLACE
+    let coords = decode_src1_coords(gid.x);
+
+    let src1_idx = params.offset_src1 + src1_idx_from_coords(coords);
+    let dst_idx = params.offset_view + view_rel_from_coords(coords);
+
+    dst[dst_idx] = src1[src1_idx];
+#else
+    let rel = select(params.ne, gid.x - params.offset_view, gid.x >= params.offset_view);
+    let coords = decode_view_coords(rel);
+
+    if (rel < params.stride_dst13 * params.src1_ne3 && in_set_view(rel, coords)) {
+        dst[gid.x] = src1[params.offset_src1 + src1_idx_from_coords(coords)];
+    } else {
+        dst[gid.x] = src0[params.offset_src0 + gid.x];
+    }
+#endif
+}

ggml/src/ggml-webgpu/wgsl-shaders/solve_tri.wgsl

@@ -0,0 +1,121 @@
+@group(0) @binding(0)
+var<storage, read_write> src0: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> src1: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read_write> dst: array<f32>;
+
+struct Params {
+    offset_src0: u32,
+    offset_src1: u32,
+    offset_dst: u32,
+
+    stride_src00: u32,
+    stride_src01: u32,
+    stride_src02: u32,
+    stride_src03: u32,
+
+    stride_src10: u32,
+    stride_src11: u32,
+    stride_src12: u32,
+    stride_src13: u32,
+
+    stride_dst0: u32,
+    stride_dst1: u32,
+    stride_dst2: u32,
+    stride_dst3: u32,
+
+    k: u32,
+    ne2: u32,
+    ne3: u32,
+};
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+var<workgroup> shA: array<f32, BATCH_N * N>;
+var<workgroup> shB: array<f32, BATCH_N * K_TILE>;
+
+fn src0_idx(row: u32, col: u32, i2: u32, i3: u32) -> u32 {
+    return params.offset_src0 +
+           col * params.stride_src00 +
+           row * params.stride_src01 +
+           i2 * params.stride_src02 +
+           i3 * params.stride_src03;
+}
+
+fn src1_idx(row: u32, col: u32, i2: u32, i3: u32) -> u32 {
+    return params.offset_src1 +
+           col * params.stride_src10 +
+           row * params.stride_src11 +
+           i2 * params.stride_src12 +
+           i3 * params.stride_src13;
+}
+
+fn dst_idx(row: u32, col: u32, i2: u32, i3: u32) -> u32 {
+    return params.offset_dst +
+           col * params.stride_dst0 +
+           row * params.stride_dst1 +
+           i2 * params.stride_dst2 +
+           i3 * params.stride_dst3;
+}
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+    @builtin(workgroup_id) workgroup_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>
+) {
+    let batch = workgroup_id.y;
+    let col = workgroup_id.x * WG_SIZE + local_id.x;
+    let i3 = batch / params.ne2;
+    let i2 = batch % params.ne2;
+    let active_lane = local_id.x < K_TILE;
+    let active_col = active_lane && col < params.k;
+
+    var X: array<f32, N>;
+
+    for (var row_base = 0u; row_base < N; row_base += BATCH_N) {
+        let cur_n = min(BATCH_N, N - row_base);
+
+        for (var i = local_id.x; i < cur_n * N; i += WG_SIZE) {
+            let tile_row = i / N;
+            let tile_col = i % N;
+            shA[i] = src0[src0_idx(row_base + tile_row, tile_col, i2, i3)];
+        }
+
+        for (var i = local_id.x; i < cur_n * K_TILE; i += WG_SIZE) {
+            let tile_row = i / K_TILE;
+            let tile_col = i % K_TILE;
+            let global_col = workgroup_id.x * WG_SIZE + tile_col;
+            let sh_idx = tile_row * K_TILE + tile_col;
+
+            if (global_col < params.k) {
+                shB[sh_idx] = src1[src1_idx(row_base + tile_row, global_col, i2, i3)];
+            } else {
+                shB[sh_idx] = 0.0;
+            }
+        }
+
+        workgroupBarrier();
+
+        if (active_col) {
+            for (var row_offset = 0u; row_offset < cur_n; row_offset++) {
+                let r = row_base + row_offset;
+                var b = shB[row_offset * K_TILE + local_id.x];
+                let a_row = row_offset * N;
+
+                for (var t = 0u; t < r; t++) {
+                    b -= shA[a_row + t] * X[t];
+                }
+
+                let x = b / shA[a_row + r];
+                X[r] = x;
+                dst[dst_idx(r, col, i2, i3)] = x;
+            }
+        }
+
+        workgroupBarrier();
+    }
+}

ggml/src/ggml-webgpu/wgsl-shaders/ssm_conv.wgsl

@@ -0,0 +1,65 @@
+@group(0) @binding(0)
+var<storage, read_write> src0: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> src1: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read_write> dst: array<f32>;
+
+struct Params {
+    offset_src0: u32,
+    offset_src1: u32,
+    offset_dst: u32,
+
+    stride_src01: u32,
+    stride_src02: u32,
+    stride_src11: u32,
+
+    stride_dst0: u32,
+    stride_dst1: u32,
+    stride_dst2: u32,
+
+    nc: u32,
+    nr: u32,
+    n_t: u32,
+    n_s: u32,
+    token_tiles: u32,
+};
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+@compute @workgroup_size(BLOCK_SIZE, TOKENS_PER_WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i1 = gid.x;
+    let tile_y = gid.y / TOKENS_PER_WG;
+    let local_token = gid.y % TOKENS_PER_WG;
+    let i3 = tile_y / params.token_tiles;
+    let token_tile = tile_y % params.token_tiles;
+    let i2 = token_tile * TOKENS_PER_WG + local_token;
+
+    if (i1 >= params.nr || i2 >= params.n_t || i3 >= params.n_s) {
+        return;
+    }
+
+    let src0_base = params.offset_src0 + i3 * params.stride_src02 + i2 + i1 * params.stride_src01;
+    let src1_base = params.offset_src1 + i1 * params.stride_src11;
+
+    var sum = 0.0;
+
+#ifdef VECTORIZED
+    sum =
+        src0[src0_base + 0u] * src1[src1_base + 0u] +
+        src0[src0_base + 1u] * src1[src1_base + 1u] +
+        src0[src0_base + 2u] * src1[src1_base + 2u] +
+        src0[src0_base + 3u] * src1[src1_base + 3u];
+#else
+    for (var i0 = 0u; i0 < params.nc; i0++) {
+        sum += src0[src0_base + i0] * src1[src1_base + i0];
+    }
+#endif
+
+    let dst_idx = params.offset_dst + i3 * params.stride_dst2 + i2 * params.stride_dst1 + i1 * params.stride_dst0;
+    dst[dst_idx] = sum;
+}