[webgpu] Optimize LinearAttention Op with subgroup (#28412)

daijh · web-flow · commit a8260cd64034 · 2026-05-12T14:28:03.000-07:00
### Description
Optimize the `LinearAttention` Op with subgroupAdd().

- Detect the adapter's `subgroupMinSize` to allocate workgroup shared
memory
- Drastically reduces workgroup shared memory usage (workgroup_size_x *
TILE_V → MAX_SG * TILE_V)
- Eliminates most `workgroupBarrier()` calls in the subgroup reduction

The optimization is gated behind `subgroup_min_size`, which is enabled
when the device supports `wgpu::FeatureName::Subgroups`. 
The original is preserved as fallback.

## Qwen3.5-4B Performance Benchmarks

| Metric | Prefill Speed (TPS) | Decode Speed (TPS) | Prefill
Improvement |
| :--- | :--- | :--- | :--- |
| **Default** | 719.6 | 29.6 | - |
| **Optimized** | 929.8 | 29.7 | 1.29x |

**Test Environment:**
* **Hardware:** Intel Panther Lake
* **Configuration:** Prefill: 1024, Decode: 128


### Motivation and Context
See above.
diff --git a/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc
@@ -81,6 +81,7 @@ Status LinearAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const {
   return WGSL_TEMPLATE_APPLY(shader, "bert/linear_attention.wgsl.template",
                              WGSL_TEMPLATE_PARAMETER(decay_broadcast_dk, decay_broadcast_dk_),
                              WGSL_TEMPLATE_PARAMETER(has_initial_state, has_initial_state_),
+                             WGSL_TEMPLATE_PARAMETER(subgroup_min_size, subgroup_min_size_),
                              WGSL_TEMPLATE_PARAMETER(tile_v, tile_v_),
                              WGSL_TEMPLATE_PARAMETER(update_rule, update_rule_int),
                              WGSL_TEMPLATE_PARAMETER(use_vec4, use_vec4));
@@ -242,7 +243,12 @@ Status LinearAttention::ComputeInternal(ComputeContext& context) const {
     }
   }
 
-  LinearAttentionProgram program{update_rule_, has_initial_state, has_decay, has_beta, decay_broadcast_dk, tile_v, components};
+  // subgroup_min_size > 0 enables subgroup-based reduction; 0 falls back to barrier-tree.
+  int subgroup_min_size = context.HasFeature(wgpu::FeatureName::Subgroups)
+                              ? static_cast<int>(context.AdapterInfo().subgroupMinSize)
+                              : 0;
+
+  LinearAttentionProgram program{update_rule_, has_initial_state, has_decay, has_beta, decay_broadcast_dk, tile_v, components, subgroup_min_size};
 
   program.AddInputs({{query, ProgramTensorMetadataDependency::TypeAndRank},
                      {key, ProgramTensorMetadataDependency::TypeAndRank},
@@ -263,7 +269,7 @@ Status LinearAttention::ComputeInternal(ComputeContext& context) const {
   program.SetDispatchGroupSize(num_workgroups)
       .SetWorkgroupSize(workgroup_size)
       .CacheHint(std::to_string(static_cast<int>(update_rule_)),
-                 has_initial_state, has_decay, has_beta, decay_broadcast_dk, tile_v, components)
+                 has_initial_state, has_decay, has_beta, decay_broadcast_dk, tile_v, components, subgroup_min_size)
       .AddUniformVariables({{static_cast<uint32_t>(batch_size)},
                             {static_cast<uint32_t>(kv_num_heads_)},
                             {static_cast<uint32_t>(seq_length)},
diff --git a/onnxruntime/contrib_ops/webgpu/bert/linear_attention.h b/onnxruntime/contrib_ops/webgpu/bert/linear_attention.h
@@ -32,15 +32,17 @@ LinearAttentionUpdateRule ParseUpdateRule(const std::string& rule_str);
 class LinearAttentionProgram final : public Program<LinearAttentionProgram> {
  public:
   LinearAttentionProgram(LinearAttentionUpdateRule update_rule, bool has_initial_state,
-                         bool has_decay, bool has_beta, bool decay_broadcast_dk, int tile_v, int components)
+                         bool has_decay, bool has_beta, bool decay_broadcast_dk, int tile_v, int components,
+                         int subgroup_min_size)
       : Program{"LinearAttention"},
         update_rule_(update_rule),
         has_initial_state_(has_initial_state),
         has_decay_(has_decay),
         has_beta_(has_beta),
         decay_broadcast_dk_(decay_broadcast_dk),
         tile_v_(tile_v),
-        components_(components) {}
+        components_(components),
+        subgroup_min_size_(subgroup_min_size) {}
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
 
@@ -65,6 +67,7 @@ class LinearAttentionProgram final : public Program<LinearAttentionProgram> {
   bool decay_broadcast_dk_;
   int tile_v_;
   int components_;
+  int subgroup_min_size_;
 };
 
 // Kernel for LinearAttention
diff --git a/onnxruntime/contrib_ops/webgpu/bert/linear_attention.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/linear_attention.wgsl.template
@@ -19,6 +19,7 @@
 #param has_initial_state
 #param decay_broadcast_dk
 #param tile_v
+#param subgroup_min_size // 0 = disable subgroup reduction
 
 // Update rule constants
 #define UPDATE_LINEAR 0
@@ -41,16 +42,29 @@ const TILE_V: u32 = tile_v;
 
 // Shared memory for parallel reduction across dk threads.
 #if update_rule == UPDATE_DELTA || update_rule == UPDATE_GATED_DELTA
+#if subgroup_min_size
+const MAX_SG: u32 = (workgroup_size_x + subgroup_min_size - 1u) / subgroup_min_size;
+var<workgroup> sg_retrieved: array<vtype, MAX_SG * TILE_V>;
+var<workgroup> sg_preout: array<vtype, MAX_SG * TILE_V>;
+var<workgroup> sg_kq: array<f32, MAX_SG>;
+var<workgroup> broadcast_buf: array<vtype, TILE_V>;
+#else
 // Fused reduction: retrieved (S^T@k), pre_output (S^T@q), and kq_dot (k^T@q)
 // are reduced in a single barrier tree.
 var<workgroup> red_retrieved: array<vtype, workgroup_size_x * TILE_V>;
 var<workgroup> red_preout: array<vtype, workgroup_size_x * TILE_V>;
 var<workgroup> red_kq: array<f32, workgroup_size_x>;
 var<workgroup> broadcast_buf: array<vtype, TILE_V>;
+#endif
+#else
+#if subgroup_min_size
+const MAX_SG: u32 = (workgroup_size_x + subgroup_min_size - 1u) / subgroup_min_size;
+var<workgroup> sg_buf: array<vtype, MAX_SG * TILE_V>;
 #else
 // Output-only reduction for linear/gated.
 var<workgroup> reduction_buf: array<vtype, workgroup_size_x * TILE_V>;
 #endif
+#endif
 
 $MAIN {
   // Identify which (batch, head, dv_tile) this workgroup handles
@@ -61,6 +75,11 @@ $MAIN {
   let dk_idx = local_idx;  // thread index = row in state matrix
   let dv_start = dv_tile_idx * TILE_V;
 
+#if subgroup_min_size
+  let sg_idx = local_idx / sg_size;
+  let num_sgs = (workgroup_size_x + sg_size - 1u) / sg_size;
+#endif
+
   // Precompute packed strides for 3D packed inputs (B, T, H*D)
   // Q: (B, T, q_num_heads * dk), K: (B, T, n_k_heads * dk),
   // V/output: (B, T, num_heads * dv)  [schema: output_dim == V_dim]
@@ -133,8 +152,47 @@ $MAIN {
     }
 
     // Fused reduction: compute retrieved = S^T@k, pre_output = S^T@q_0,
-    // and kq_dot = k^T@q_0 in a single barrier tree.
+    // and kq_dot = k^T@q_0.
     // Then: output_0 = scale * (pre_output + delta * kq_dot)
+#if subgroup_min_size
+    for (var j = 0u; j < TILE_V; j++) {
+      let ret_j = subgroupAdd(state[j] * k_val);
+      let pre_j = subgroupAdd(state[j] * q0_val);
+      if (sg_id == 0u) {
+        sg_retrieved[sg_idx * TILE_V + j] = ret_j;
+        sg_preout[sg_idx * TILE_V + j] = pre_j;
+      }
+    }
+    let kq_sg = subgroupAdd(k_val * q0_val);
+    if (sg_id == 0u) {
+      sg_kq[sg_idx] = kq_sg;
+    }
+    workgroupBarrier();
+
+    // Threads 0..TILE_V-1 each handle one dv position independently.
+    let v_base = bt_offset * packed_dv + head_idx * uniforms.head_dim_v + dv_start;
+    let beta_base = bt_offset * uniforms.num_heads + head_idx;
+    if (dk_idx < TILE_V) {
+      let j = dk_idx;
+      var retrieved_sum = vtype(0.0);
+      var preout_sum = vtype(0.0);
+      var kq_dot: f32 = 0.0;
+      for (var sg = 0u; sg < num_sgs; sg++) {
+        retrieved_sum += sg_retrieved[sg * TILE_V + j];
+        preout_sum += sg_preout[sg * TILE_V + j];
+        kq_dot += sg_kq[sg];
+      }
+      let beta_val = f32(beta[beta_base]);
+      if (dv_start + j < uniforms.head_dim_v) {
+        let v_val = vtype(value[v_base + j]);
+        let delta_j = beta_val * (v_val - retrieved_sum);
+        broadcast_buf[j] = delta_j;
+        output[bt_offset * packed_dv + out_head_0 * uniforms.head_dim_v + dv_start + j] = otype((preout_sum + delta_j * kq_dot) * uniforms.scale);
+      } else {
+        broadcast_buf[j] = vtype(0.0);
+      }
+    }
+#else
     for (var j = 0u; j < TILE_V; j++) {
       red_retrieved[j * workgroup_size_x + dk_idx] = state[j] * k_val;
       red_preout[j * workgroup_size_x + dk_idx] = state[j] * q0_val;
@@ -172,6 +230,7 @@ $MAIN {
         }
       }
     }
+#endif
     workgroupBarrier();
 
     // All threads: update state with delta (S_new = S_old + k * delta)
@@ -188,6 +247,27 @@ $MAIN {
       if (dk_idx < uniforms.head_dim_k) {
         qg_val = f32(query[bt_offset * packed_dk_q + q_head_g * uniforms.head_dim_k + dk_idx]);
       }
+#if subgroup_min_size
+      // Subgroup reduction of state*qg
+      for (var j = 0u; j < TILE_V; j++) {
+        let pre_j = subgroupAdd(state[j] * qg_val);
+        if (sg_id == 0u) {
+          sg_preout[sg_idx * TILE_V + j] = pre_j;
+        }
+      }
+      workgroupBarrier();
+      // Cross-subgroup reduction + parallel output write
+      if (dk_idx < TILE_V) {
+        let j = dk_idx;
+        var preout_sum = vtype(0.0);
+        for (var sg = 0u; sg < num_sgs; sg++) {
+          preout_sum += sg_preout[sg * TILE_V + j];
+        }
+        if (dv_start + j < uniforms.head_dim_v) {
+          output[bt_offset * packed_dv + q_head_g * uniforms.head_dim_v + dv_start + j] = otype(preout_sum * uniforms.scale);
+        }
+      }
+#else
       for (var j = 0u; j < TILE_V; j++) {
         red_preout[j * workgroup_size_x + dk_idx] = state[j] * qg_val;
       }
@@ -207,6 +287,7 @@ $MAIN {
           }
         }
       }
+#endif
       workgroupBarrier();
     }
 
@@ -228,6 +309,26 @@ $MAIN {
         if (dk_idx < uniforms.head_dim_k) {
           q_val_g = f32(query[bt_offset * packed_dk_q + q_head_idx * uniforms.head_dim_k + dk_idx]);
         }
+#if subgroup_min_size
+        for (var j = 0u; j < TILE_V; j++) {
+          let pre_j = subgroupAdd(state[j] * q_val_g);
+          if (sg_id == 0u) {
+            sg_buf[sg_idx * TILE_V + j] = pre_j;
+          }
+        }
+        workgroupBarrier();
+        if (dk_idx < TILE_V) {
+          let j = dk_idx;
+          var preout_sum = vtype(0.0);
+          for (var sg = 0u; sg < num_sgs; sg++) {
+            preout_sum += sg_buf[sg * TILE_V + j];
+          }
+          let out_base = bt_offset * packed_dv + q_head_idx * uniforms.head_dim_v + dv_start;
+          if (dv_start + j < uniforms.head_dim_v) {
+            output[out_base + j] = otype(preout_sum * uniforms.scale);
+          }
+        }
+#else
         for (var j = 0u; j < TILE_V; j++) {
           reduction_buf[j * workgroup_size_x + dk_idx] = state[j] * q_val_g;
         }
@@ -248,6 +349,7 @@ $MAIN {
             }
           }
         }
+#endif
         workgroupBarrier();
       }
     } else {
@@ -257,6 +359,26 @@ $MAIN {
       if (dk_idx < uniforms.head_dim_k) {
         q_val_inv = f32(query[bt_offset * packed_dk_q + q_head_inv * uniforms.head_dim_k + dk_idx]);
       }
+#if subgroup_min_size
+      for (var j = 0u; j < TILE_V; j++) {
+        let pre_j = subgroupAdd(state[j] * q_val_inv);
+        if (sg_id == 0u) {
+          sg_buf[sg_idx * TILE_V + j] = pre_j;
+        }
+      }
+      workgroupBarrier();
+      if (dk_idx < TILE_V) {
+        let j = dk_idx;
+        var preout_sum = vtype(0.0);
+        for (var sg = 0u; sg < num_sgs; sg++) {
+          preout_sum += sg_buf[sg * TILE_V + j];
+        }
+        let out_base = bt_offset * packed_dv + head_idx * uniforms.head_dim_v + dv_start;
+        if (dv_start + j < uniforms.head_dim_v) {
+          output[out_base + j] = otype(preout_sum * uniforms.scale);
+        }
+      }
+#else
       for (var j = 0u; j < TILE_V; j++) {
         reduction_buf[j * workgroup_size_x + dk_idx] = state[j] * q_val_inv;
       }
@@ -277,6 +399,7 @@ $MAIN {
           }
         }
       }
+#endif
       workgroupBarrier();
     }
 #endif
