[ExecuTorch][WebGPU] Optimize Qwen3 long-context FlashDecoding

backends/webgpu/runtime/WebGPUShaderRegistry.cpp

@@ -119,7 +119,9 @@
 #include <executorch/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/select/select_wgsl.h>
@@ -155,7 +157,7 @@
 namespace executorch::backends::webgpu {
 namespace {
 
-constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
+constexpr std::array<WebGPUShaderInfo, 139> kShaderRegistry = {{
     {
         "abs",
         kAbsWGSL,
@@ -961,13 +963,27 @@ constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
         kSdpaFdReduceWorkgroupSizeY,
         kSdpaFdReduceWorkgroupSizeZ,
     },
+    {
+        "sdpa_fd_reduce_gqa2_f16",
+        kSdpaFdReduceGqa2F16WGSL,
+        kSdpaFdReduceGqa2F16WorkgroupSizeX,
+        kSdpaFdReduceGqa2F16WorkgroupSizeY,
+        kSdpaFdReduceGqa2F16WorkgroupSizeZ,
+    },
     {
         "sdpa_fd_split",
         kSdpaFdSplitWGSL,
         kSdpaFdSplitWorkgroupSizeX,
         kSdpaFdSplitWorkgroupSizeY,
         kSdpaFdSplitWorkgroupSizeZ,
     },
+    {
+        "sdpa_fd_split_gqa2_f16",
+        kSdpaFdSplitGqa2F16WGSL,
+        kSdpaFdSplitGqa2F16WorkgroupSizeX,
+        kSdpaFdSplitGqa2F16WorkgroupSizeY,
+        kSdpaFdSplitGqa2F16WorkgroupSizeZ,
+    },
     {
         "sdpa_fd_split_half",
         kSdpaFdSplitHalfWGSL,

backends/webgpu/runtime/ops/sdpa/Sdpa.cpp

@@ -694,7 +694,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     // lambda runs, so eager construction here can never throw on it.
     if (fd_eligible) {
       state.fd = make_sdpa_fd_decode_state(
-          gr.device(), Hq, Hkv, D, state.context_len, g, scale);
+          gr.device(), Hq, Hkv, D, state.context_len, g, scale, gr.kv_f16());
     }
     return state;
   };

backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp

@@ -11,7 +11,9 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h>
+#include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h>
 
@@ -27,14 +29,23 @@ namespace executorch::backends::webgpu {
 
 namespace {
 
-// MUST match the .wgsl: MAX_SPLITS and WG_SIZE*MAX_D_PER_LANE.
-constexpr uint32_t kSdpaFdSplitTile = 64; // KV positions per split
-constexpr uint32_t kSdpaFdMaxSplits = 128; // == MAX_SPLITS in both .wgsl files
 // Public head-dim limit (kSdpaFdMaxHeadDim) must equal the kernel's lane-owns-D
 // reach; tie them so a WG_SIZE change can't silently desync the Sdpa.cpp gate.
 static_assert(
     kSdpaFdMaxHeadDim == kSdpaFdSplitWorkgroupSizeX * 2u,
     "kSdpaFdMaxHeadDim must match WG_SIZE * MAX_D_PER_LANE");
+static_assert(
+    kSdpaFdSplitGqa2F16WorkgroupSizeX == kSdpaFdSplitWorkgroupSizeX,
+    "Qwen GQA2 and generic split workgroup widths must match");
+static_assert(
+    kSdpaFdReduceGqa2F16WorkgroupSizeX == kSdpaFdReduceWorkgroupSizeX,
+    "Qwen GQA2 and generic reduce workgroup widths must match");
+static_assert(
+    kSdpaFdMaxHeadDim == kSdpaFdSplitGqa2F16WorkgroupSizeX * 2u,
+    "Qwen GQA2 split must reach D=128");
+static_assert(
+    kSdpaFdMaxSplits == 128u,
+    "MAX_SPLITS must match generic and GQA2 split/reduce shaders");
 
 struct FdSplitParams {
   uint32_t _pad0; // 16B-alignment pad (head index derived from workgroup_id)
@@ -138,7 +149,8 @@ SdpaFdDecodeState make_sdpa_fd_decode_state(
     int64_t D,
     int64_t context_len,
     int64_t g,
-    float scale) {
+    float scale,
+    bool kv_f16) {
   if (Hq <= 0 || Hkv <= 0 || D <= 0 || context_len <= 0 || g <= 0) {
     throw std::runtime_error(
         "WebGPU sdpa FlashDecoding: dimensions must be positive");
@@ -162,17 +174,24 @@ SdpaFdDecodeState make_sdpa_fd_decode_state(
         "WebGPU sdpa FlashDecoding: head dim must be a multiple of 4");
   }
 
-  uint32_t num_splits = static_cast<uint32_t>(
-      (context_len + kSdpaFdSplitTile - 1) / kSdpaFdSplitTile);
-  num_splits = std::min(num_splits, kSdpaFdMaxSplits);
+  const bool qwen_gqa2_f16 = is_qwen_gqa2_f16_fd_route(kv_f16, Hq, Hkv, D, g);
+  const uint32_t num_splits =
+      sdpa_fd_num_splits(static_cast<uint32_t>(context_len), qwen_gqa2_f16);
   const uint32_t split_len =
       static_cast<uint32_t>((context_len + num_splits - 1) / num_splits);
 
-  const uint64_t split_threads = static_cast<uint64_t>(Hq) *
+  const uint32_t split_heads = sdpa_fd_split_head_count(Hq, Hkv, qwen_gqa2_f16);
+  const uint32_t split_workgroup_size = qwen_gqa2_f16
+      ? kSdpaFdSplitGqa2F16WorkgroupSizeX
+      : kSdpaFdSplitWorkgroupSizeX;
+  const uint32_t reduce_workgroup_size = qwen_gqa2_f16
+      ? kSdpaFdReduceGqa2F16WorkgroupSizeX
+      : kSdpaFdReduceWorkgroupSizeX;
+  const uint64_t split_threads = static_cast<uint64_t>(split_heads) *
       static_cast<uint64_t>(num_splits) *
-      static_cast<uint64_t>(kSdpaFdSplitWorkgroupSizeX);
+      static_cast<uint64_t>(split_workgroup_size);
   const uint64_t reduce_threads =
-      static_cast<uint64_t>(Hq) * kSdpaFdReduceWorkgroupSizeX;
+      static_cast<uint64_t>(Hq) * reduce_workgroup_size;
   if (split_threads > UINT32_MAX || reduce_threads > UINT32_MAX) {
     throw std::runtime_error(
         "WebGPU sdpa FlashDecoding: thread count exceeds uint32 max");
@@ -181,12 +200,12 @@ SdpaFdDecodeState make_sdpa_fd_decode_state(
   const uint32_t split_wgc = utils::compute_1d_workgroup_count(
       device,
       static_cast<uint32_t>(split_threads),
-      kSdpaFdSplitWorkgroupSizeX,
+      split_workgroup_size,
       "fd_split");
   const uint32_t reduce_wgc = utils::compute_1d_workgroup_count(
       device,
       static_cast<uint32_t>(reduce_threads),
-      kSdpaFdReduceWorkgroupSizeX,
+      reduce_workgroup_size,
       "fd_reduce");
   return {
       static_cast<uint32_t>(Hq),
@@ -197,6 +216,7 @@ SdpaFdDecodeState make_sdpa_fd_decode_state(
       num_splits,
       split_len,
       scale,
+      qwen_gqa2_f16,
       {split_wgc, 1u},
       {reduce_wgc, 1u}};
 }
@@ -230,7 +250,11 @@ SdpaFdDecodeResources record_sdpa_fd_decode_dispatches(
       {k_cache.buffer, k_cache.nbytes},
       {v_cache.buffer, v_cache.nbytes}};
   const char* split_shader = kSdpaFdSplitWGSL;
-  if (graph.kv_f16()) {
+  const char* split_label = "fd_split";
+  if (state.qwen_gqa2_f16) {
+    split_shader = kSdpaFdSplitGqa2F16WGSL;
+    split_label = "fd_split_gqa2_f16";
+  } else if (graph.kv_f16()) {
     split_shader = kSdpaFdSplitHalfWGSL;
   }
   build_dispatch(
@@ -243,7 +267,7 @@ SdpaFdDecodeResources record_sdpa_fd_decode_dispatches(
       sizeof(sp),
       state.split_grid.x,
       true,
-      "fd_split");
+      split_label);
 
   // Pass 2: reduce (Hq WGs) -> reads part_o, part_ml; writes out.
   FdReduceParams rp = make_reduce_params(state);
@@ -252,17 +276,21 @@ SdpaFdDecodeResources record_sdpa_fd_decode_dispatches(
       {out.buffer, out.nbytes},
       {part_o, po_floats * sizeof(float)},
       {part_ml, pml_floats * sizeof(float)}};
+  const char* reduce_shader =
+      state.qwen_gqa2_f16 ? kSdpaFdReduceGqa2F16WGSL : kSdpaFdReduceWGSL;
+  const char* reduce_label =
+      state.qwen_gqa2_f16 ? "fd_reduce_gqa2_f16" : "fd_reduce";
   build_dispatch(
       graph,
-      kSdpaFdReduceWGSL,
+      reduce_shader,
       reduce_bindings,
       3,
       1,
       ub_reduce,
       sizeof(rp),
       state.reduce_grid.x,
       true,
-      "fd_reduce");
+      reduce_label);
 
   return {ub_split, ub_reduce, {dispatch_begin, graph.num_dispatches()}};
 }

backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h

@@ -19,6 +19,35 @@ namespace executorch::backends::webgpu {
 // MAX_D_PER_LANE(2). Decode shapes above this fall through to the materialized
 // SDPA path (the FD selection predicate in Sdpa.cpp checks this).
 constexpr int64_t kSdpaFdMaxHeadDim = 128;
+constexpr uint32_t kSdpaFdGenericSplitTile = 64u;
+constexpr uint32_t kSdpaFdQwenGqa2SplitTile = 128u;
+constexpr uint32_t kSdpaFdMaxSplits = 128u;
+
+constexpr bool is_qwen_gqa2_f16_fd_route(
+    bool kv_f16,
+    int64_t hq,
+    int64_t hkv,
+    int64_t d,
+    int64_t g) {
+  return kv_f16 && hq == 16 && hkv == 8 && d == 128 && g == 2;
+}
+
+constexpr uint32_t sdpa_fd_split_tile(bool qwen_gqa2_f16) {
+  return qwen_gqa2_f16 ? kSdpaFdQwenGqa2SplitTile : kSdpaFdGenericSplitTile;
+}
+
+constexpr uint32_t sdpa_fd_num_splits(
+    uint32_t context_len,
+    bool qwen_gqa2_f16) {
+  const uint32_t tile = sdpa_fd_split_tile(qwen_gqa2_f16);
+  const uint32_t splits = context_len / tile + (context_len % tile != 0u);
+  return splits > kSdpaFdMaxSplits ? kSdpaFdMaxSplits : splits;
+}
+
+constexpr uint32_t
+sdpa_fd_split_head_count(int64_t hq, int64_t hkv, bool qwen_gqa2_f16) {
+  return static_cast<uint32_t>(qwen_gqa2_f16 ? hkv : hq);
+}
 
 struct SdpaFdDecodeState {
   uint32_t Hq;
@@ -29,6 +58,7 @@ struct SdpaFdDecodeState {
   uint32_t num_splits;
   uint32_t split_len;
   float scale;
+  bool qwen_gqa2_f16;
   utils::WgCount split_grid;
   utils::WgCount reduce_grid;
 };
@@ -46,7 +76,8 @@ SdpaFdDecodeState make_sdpa_fd_decode_state(
     int64_t D,
     int64_t context_len,
     int64_t g,
-    float scale);
+    float scale,
+    bool kv_f16);
 
 // Records split + reduce with retained UBOs. Route selection is owned by the
 // caller so this helper never mutates recorded dispatch counts.

backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16.wgsl

@@ -0,0 +1,83 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
+@group(0) @binding(1) var<storage, read> t_part_o: array<f32>;
+@group(0) @binding(2) var<storage, read> t_part_ml: array<f32>;
+
+struct Params {
+  D: u32,
+  num_splits: u32,
+  _pad0: u32,
+  _pad1: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+const WG_SIZE: u32 = 64u;
+const MAX_SPLITS: u32 = 128u;
+const MAX_D_PER_LANE: u32 = 2u;
+const NEG_INF: f32 = -1.0e30;
+
+// w_i = exp(m_i - M) per split, computed once and reused for the L-sum and every output dim.
+var<workgroup> sh_w: array<f32, MAX_SPLITS>;
+var<workgroup> sh_reduce: array<f32, WG_SIZE>;
+
+// FlashDecoding pass 2: online-softmax merge of the per-split partials, then normalize.
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  let h = wid.x;
+  let t = lid.x;
+  let D = params.D;
+  let ns = params.num_splits;
+  let head_base = h * MAX_SPLITS;
+
+  var local_m: f32 = NEG_INF;
+  for (var i: u32 = t; i < ns; i = i + WG_SIZE) {
+    local_m = max(local_m, t_part_ml[(head_base + i) * 2u + 0u]);
+  }
+  sh_reduce[t] = local_m;
+  workgroupBarrier();
+  for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
+    if (t < stride) {
+      sh_reduce[t] = max(sh_reduce[t], sh_reduce[t + stride]);
+    }
+    workgroupBarrier();
+  }
+  let M = sh_reduce[0];
+  // Compute w_i = exp(m_i - M) once per split into shared memory (was recomputed per output dim).
+  for (var i: u32 = t; i < ns; i = i + WG_SIZE) {
+    sh_w[i] = exp(t_part_ml[(head_base + i) * 2u + 0u] - M);
+  }
+  workgroupBarrier();
+
+  var local_l: f32 = 0.0;
+  for (var i: u32 = t; i < ns; i = i + WG_SIZE) {
+    local_l = local_l + sh_w[i] * t_part_ml[(head_base + i) * 2u + 1u];
+  }
+  sh_reduce[t] = local_l;
+  workgroupBarrier();
+  for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
+    if (t < stride) {
+      sh_reduce[t] = sh_reduce[t] + sh_reduce[t + stride];
+    }
+    workgroupBarrier();
+  }
+  let L = sh_reduce[0];
+  let inv = select(0.0, 1.0 / L, L > 0.0);
+
+  for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) {
+    let d = t + nd * WG_SIZE;
+    if (d < D) {
+      var acc: f32 = 0.0;
+      for (var i: u32 = 0u; i < ns; i = i + 1u) {
+        acc = acc + sh_w[i] * t_part_o[(head_base + i) * D + d];
+      }
+      t_out[h * D + d] = acc * inv;
+    }
+  }
+}