diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py
index 68399681699..ffbbc796c62 100644
--- a/backends/vulkan/custom_ops_lib.py
+++ b/backends/vulkan/custom_ops_lib.py
@@ -900,6 +900,44 @@ def apply_rotary_emb_hf_meta(
 lib.impl(name, apply_rotary_emb_hf_meta, "Meta")
 apply_rotary_emb_hf_op = getattr(getattr(torch.ops, namespace), name)
 
+################################
+## apply_rotary_emb_hf_single ##
+################################
+
+
+def apply_rotary_emb_hf_single_impl(
+    x: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+    start_pos: int,
+):
+    seq_len = x.shape[1]
+    freqs_cos = freqs_cos[start_pos : start_pos + seq_len]
+    freqs_sin = freqs_sin[start_pos : start_pos + seq_len]
+    pattern = vk_patterns.HfRotaryEmbeddingSinglePattern()
+    return pattern.forward(x, freqs_cos, freqs_sin)
+
+
+def apply_rotary_emb_hf_single_meta(
+    x: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+    start_pos: int,
+):
+    output_dtype = torch.promote_types(
+        torch.promote_types(x.dtype, freqs_cos.dtype), freqs_sin.dtype
+    )
+    return torch.empty_like(x, dtype=output_dtype)
+
+
+name = "apply_rotary_emb_hf_single"
+lib.define(
+    f"{name}(Tensor x, Tensor freqs_cos, Tensor freqs_sin, SymInt start_pos) -> Tensor"
+)
+lib.impl(name, apply_rotary_emb_hf_single_impl, "CompositeExplicitAutograd")
+lib.impl(name, apply_rotary_emb_hf_single_meta, "Meta")
+apply_rotary_emb_hf_single_op = getattr(getattr(torch.ops, namespace), name)
+
 ##################################
 ## apply_rotary_emb_interleaved ##
 ##################################
@@ -1149,6 +1187,66 @@ def sdpa_impl(
 lib.impl(name, sdpa_impl, "CompositeExplicitAutograd")
 sdpa_op = getattr(getattr(torch.ops, namespace), name)
 
+#################
+## gemma4_sdpa ##
+#################
+
+
+def gemma4_sdpa_impl(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    start_pos: int,
+    attn_mask: torch.Tensor,
+    dropout_p: float,
+    is_causal: bool,
+    scale: float,
+) -> torch.Tensor:
+    del start_pos
+    if dropout_p != 0.0 or is_causal or scale != 1.0:
+        raise ValueError("gemma4_sdpa requires dropout=0, causal=false, scale=1")
+    if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:
+        raise ValueError("gemma4_sdpa requires BSHD query, key, and value")
+    if key.shape != value.shape or query.shape[0] != key.shape[0]:
+        raise ValueError("gemma4_sdpa query, key, and value shapes do not match")
+    if query.shape[-1] != key.shape[-1] or query.shape[2] % key.shape[2] != 0:
+        raise ValueError("gemma4_sdpa requires grouped-query compatible heads")
+    if attn_mask.dim() != 2 or tuple(attn_mask.shape) != (
+        query.shape[1],
+        key.shape[1],
+    ):
+        raise ValueError("gemma4_sdpa requires a rank-2 [S_q, S_kv] mask")
+
+    group_size = query.shape[2] // key.shape[2]
+    query_bhsd = query.transpose(1, 2)
+    key_bhsd = key.transpose(1, 2).repeat_interleave(group_size, dim=1)
+    value_bhsd = value.transpose(1, 2).repeat_interleave(group_size, dim=1)
+    scores = torch.matmul(query_bhsd, key_bhsd.transpose(-2, -1))
+    scores = scores + attn_mask
+    return torch.matmul(torch.softmax(scores, dim=-1), value_bhsd).transpose(1, 2)
+
+
+def gemma4_sdpa_meta(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    start_pos: int,
+    attn_mask: torch.Tensor,
+    dropout_p: float,
+    is_causal: bool,
+    scale: float,
+) -> torch.Tensor:
+    return torch.empty_like(query)
+
+
+name = "gemma4_sdpa"
+lib.define(
+    f"{name}(Tensor query, Tensor key, Tensor value, SymInt start_pos, Tensor attn_mask, float dropout_p, bool is_causal, float scale) -> Tensor"
+)
+lib.impl(name, gemma4_sdpa_impl, "CompositeExplicitAutograd")
+lib.impl(name, gemma4_sdpa_meta, "Meta")
+gemma4_sdpa_op = getattr(getattr(torch.ops, namespace), name)
+
 ################
 ## rms_norm ##
 ################
diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h
index 60638b499bb..7884a0f02cb 100644
--- a/backends/webgpu/runtime/WebGPUDispatchMath.h
+++ b/backends/webgpu/runtime/WebGPUDispatchMath.h
@@ -101,6 +101,34 @@ constexpr bool should_record_sdpa_dual_route(
   return fd_eligible && (has_dynamic_sequence || has_dynamic_position);
 }
 
+constexpr uint32_t kCqpQdqFusedInvocations = 256u;
+constexpr uint32_t kCqpQdqFusedStorageBytes =
+    2u * kCqpQdqFusedInvocations * sizeof(float);
+
+constexpr bool is_cqp_qdq_fusion_eligible(
+    uint32_t rows,
+    uint32_t row_width,
+    uint64_t numel,
+    int64_t quant_min,
+    int64_t quant_max,
+    bool asymmetric,
+    bool per_row_block,
+    bool keepdim,
+    uint32_t max_invocations,
+    uint32_t max_workgroup_size_x,
+    uint32_t max_workgroup_storage_bytes) {
+  return rows > 0u && row_width > 0u &&
+      numel == static_cast<uint64_t>(rows) * row_width && asymmetric &&
+      per_row_block && !keepdim && quant_min == -128 && quant_max == 127 &&
+      max_invocations >= kCqpQdqFusedInvocations &&
+      max_workgroup_size_x >= kCqpQdqFusedInvocations &&
+      max_workgroup_storage_bytes >= kCqpQdqFusedStorageBytes;
+}
+
+constexpr uint32_t cqp_resize_workgroups(bool producer_elided, uint32_t grid) {
+  return producer_elided ? 0u : grid;
+}
+
 constexpr bool is_q4gsw_bk64_eligible(
     uint32_t k,
     uint32_t n,
diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp
index b7332a1cd9c..1ae8acfa578 100644
--- a/backends/webgpu/runtime/WebGPUGraph.cpp
+++ b/backends/webgpu/runtime/WebGPUGraph.cpp
@@ -1006,12 +1006,55 @@ WebGPUGraph::~WebGPUGraph() {
   }
 }
 
+void WebGPUGraph::offer_cqp_fusion_site(CqpFusionSite site) {
+  site.valid = true;
+  cqp_fusion_site_ = std::move(site);
+}
+
+WebGPUGraph::CqpFusionSite WebGPUGraph::claim_cqp_fusion_site(
+    int input_id,
+    int scales_id,
+    int zero_points_id,
+    uint32_t rows,
+    uint32_t row_width) {
+  const CqpFusionSite& site = cqp_fusion_site_;
+  const bool producer_is_choose_qparams =
+      site.dispatch_index < dispatches_.size() &&
+      dispatches_[site.dispatch_index].kernel_name == "choose_qparams_affine";
+  const bool matches = site.valid && site.input_id == input_id &&
+      site.scales_id == scales_id && site.zero_points_id == zero_points_id &&
+      site.rows == rows && site.row_width == row_width &&
+      site.input_buffer == get_tensor(input_id).buffer &&
+      site.scales_buffer == get_tensor(scales_id).buffer &&
+      site.zero_points_buffer == get_tensor(zero_points_id).buffer &&
+      site.dispatch_index + 1u == dispatches_.size() &&
+      producer_is_choose_qparams && site.producer_elided != nullptr;
+  if (!matches) {
+    return CqpFusionSite{};
+  }
+  CqpFusionSite claimed = site;
+  cqp_fusion_site_ = CqpFusionSite{};
+  return claimed;
+}
+
 void WebGPUGraph::build(
     const void* flatbuffer_data,
     const uint8_t* constant_data,
     size_t constant_data_size,
     const executorch::runtime::NamedDataMap* named_data_map,
     WebGPUGraphConfig config) {
+  clear_cqp_fusion_site();
+  clear_rms_fusion_site();
+  clear_slice_chain();
+  struct ClearFusionSitesOnExit {
+    WebGPUGraph* graph;
+    ~ClearFusionSitesOnExit() {
+      graph->clear_cqp_fusion_site();
+      graph->clear_rms_fusion_site();
+      graph->clear_slice_chain();
+    }
+  } clear_fusion_sites_on_exit{this};
+
   if (!device_) {
     auto* ctx = get_default_webgpu_context();
     if (ctx) {
diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h
index d639d0851e6..4f1453d9164 100644
--- a/backends/webgpu/runtime/WebGPUGraph.h
+++ b/backends/webgpu/runtime/WebGPUGraph.h
@@ -12,6 +12,7 @@
 
 #include <cstdint>
 #include <functional>
+#include <memory>
 #include <stdexcept>
 #include <string>
 #include <type_traits>
@@ -147,6 +148,36 @@ struct WebGPUGraphConfig {
 
 class WebGPUGraph {
  public:
+  struct CqpFusionSite {
+    bool valid = false;
+    int input_id = -1;
+    int scales_id = -1;
+    int zero_points_id = -1;
+    uint32_t rows = 0u;
+    uint32_t row_width = 0u;
+    int64_t quant_min = 0;
+    int64_t quant_max = 0;
+    size_t dispatch_index = 0u;
+    WGPUBuffer input_buffer = nullptr;
+    WGPUBuffer scales_buffer = nullptr;
+    WGPUBuffer zero_points_buffer = nullptr;
+    std::shared_ptr<bool> producer_elided;
+  };
+
+  struct RmsFusionSite {
+    bool valid = false;
+    bool add_fused = false;
+    int in_id = -1;
+    int weight_id = -1;
+    int out_id = -1;
+    int resid_id = -1;
+    int addout_id = -1;
+    uint32_t num_rows = 0u;
+    uint32_t row_width = 0u;
+    size_t dispatch_index = 0u;
+    WGPUBuffer params_buffer = nullptr;
+  };
+
   WebGPUGraph();
   ~WebGPUGraph();
 
@@ -410,6 +441,55 @@ class WebGPUGraph {
     return dispatches_.size();
   }
 
+  void offer_cqp_fusion_site(CqpFusionSite site);
+  CqpFusionSite claim_cqp_fusion_site(
+      int input_id,
+      int scales_id,
+      int zero_points_id,
+      uint32_t rows,
+      uint32_t row_width);
+  void clear_cqp_fusion_site() {
+    cqp_fusion_site_ = CqpFusionSite{};
+  }
+
+  void offer_rms_fusion_site(RmsFusionSite site) {
+    site.valid = true;
+    rms_fusion_site_ = std::move(site);
+  }
+  const RmsFusionSite& rms_fusion_site() const {
+    return rms_fusion_site_;
+  }
+  void clear_rms_fusion_site() {
+    rms_fusion_site_ = RmsFusionSite{};
+  }
+
+  // Dual-store slice merge: the preceding slice offers its dispatch so a
+  // following whole-extent copy can re-bind it to a second destination.
+  // Graph-instance state, so two graphs can never observe each other's
+  // dispatch indices or buffer handles.
+  struct SliceChain {
+    bool valid = false;
+    int out_id = -1;
+    size_t dispatch_idx = 0;
+    WGPUBuffer in_buffer = nullptr;
+    size_t in_nbytes = 0;
+    WGPUBuffer out_buffer = nullptr;
+    size_t out_nbytes = 0;
+    WGPUBuffer out_meta_buf = nullptr;
+    WGPUBuffer in_meta_buf = nullptr;
+    WGPUBuffer params_buf = nullptr;
+  };
+
+  void offer_slice_chain(SliceChain chain) {
+    slice_chain_ = chain;
+  }
+  const SliceChain& slice_chain() const {
+    return slice_chain_;
+  }
+  void clear_slice_chain() {
+    slice_chain_ = SliceChain{};
+  }
+
   size_t register_dispatch_route_group(
       const std::vector<utils::DispatchRange>& ranges) {
     validate_dynamic_dispatch_route_ranges(ranges);
@@ -761,6 +841,9 @@ class WebGPUGraph {
 
   std::vector<WebGPUDispatch> dispatches_;
   utils::DispatchRouteRegistry dispatch_routes_;
+  CqpFusionSite cqp_fusion_site_;
+  RmsFusionSite rms_fusion_site_;
+  SliceChain slice_chain_;
 
   // Prepack-routed constant sources (offset/named-key + size); the prepack node
   // materializes these once. constant_data_/named_data_map_ point at the .pte
diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp
index ebc3dcc6cd7..fb40eeeb760 100644
--- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp
+++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp
@@ -50,6 +50,7 @@
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h>
@@ -91,6 +92,7 @@
 #include <executorch/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantize/quantize_per_tensor_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/linear_dW_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h>
@@ -103,9 +105,12 @@
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/relu/relu_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved_wgsl.h>
@@ -124,6 +129,7 @@
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/select/select_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/slice/slice_dual_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/slice/slice_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/softmax/softmax_wgsl.h>
@@ -155,7 +161,7 @@
 namespace executorch::backends::webgpu {
 namespace {
 
-constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
+constexpr std::array<WebGPUShaderInfo, 143> kShaderRegistry = {{
     {
         "abs",
         kAbsWGSL,
@@ -317,6 +323,13 @@ constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
         kChooseQparamsAffineWorkgroupSizeY,
         kChooseQparamsAffineWorkgroupSizeZ,
     },
+    {
+        "choose_qparams_dq8ca_fused",
+        kChooseQparamsDq8caFusedWGSL,
+        kChooseQparamsDq8caFusedWorkgroupSizeX,
+        kChooseQparamsDq8caFusedWorkgroupSizeY,
+        kChooseQparamsDq8caFusedWorkgroupSizeZ,
+    },
     {
         "clamp",
         kClampWGSL,
@@ -485,6 +498,13 @@ constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
         kEtVkSdpaQkEntryWorkgroupSizeY,
         kEtVkSdpaQkEntryWorkgroupSizeZ,
     },
+    {
+        "et_vk_sdpa_qk_entry_exact",
+        kEtVkSdpaQkEntryExactWGSL,
+        kEtVkSdpaQkEntryExactWorkgroupSizeX,
+        kEtVkSdpaQkEntryExactWorkgroupSizeY,
+        kEtVkSdpaQkEntryExactWorkgroupSizeZ,
+    },
     {
         "exp",
         kExpWGSL,
@@ -856,6 +876,13 @@ constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
         kQcs4wLinearWorkgroupSizeY,
         kQcs4wLinearWorkgroupSizeZ,
     },
+    {
+        "quantize_dequantize_per_row",
+        kQuantizeDequantizePerRowWGSL,
+        kQuantizeDequantizePerRowWorkgroupSizeX,
+        kQuantizeDequantizePerRowWorkgroupSizeY,
+        kQuantizeDequantizePerRowWorkgroupSizeZ,
+    },
     {
         "quantize_per_tensor",
         kQuantizePerTensorWGSL,
@@ -898,6 +925,20 @@ constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
         kRmsNormVec4WorkgroupSizeY,
         kRmsNormVec4WorkgroupSizeZ,
     },
+    {
+        "rms_norm_vec4_add",
+        kRmsNormVec4AddWGSL,
+        kRmsNormVec4AddWorkgroupSizeX,
+        kRmsNormVec4AddWorkgroupSizeY,
+        kRmsNormVec4AddWorkgroupSizeZ,
+    },
+    {
+        "rms_norm_vec4_add_scale",
+        kRmsNormVec4AddScaleWGSL,
+        kRmsNormVec4AddScaleWorkgroupSizeX,
+        kRmsNormVec4AddScaleWorkgroupSizeY,
+        kRmsNormVec4AddScaleWorkgroupSizeZ,
+    },
     {
         "rotary_embedding",
         kRotaryEmbeddingWGSL,
@@ -1017,6 +1058,13 @@ constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
         kSliceWorkgroupSizeY,
         kSliceWorkgroupSizeZ,
     },
+    {
+        "slice_dual",
+        kSliceDualWGSL,
+        kSliceDualWorkgroupSizeX,
+        kSliceDualWorkgroupSizeY,
+        kSliceDualWorkgroupSizeZ,
+    },
     {
         "softmax",
         kSoftmaxWGSL,
diff --git a/backends/webgpu/runtime/ops/add/BinaryOp.cpp b/backends/webgpu/runtime/ops/add/BinaryOp.cpp
index 8a8f10302f9..217b2f72f94 100644
--- a/backends/webgpu/runtime/ops/add/BinaryOp.cpp
+++ b/backends/webgpu/runtime/ops/add/BinaryOp.cpp
@@ -11,6 +11,7 @@
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/TensorMeta.h>
 #include <executorch/backends/webgpu/runtime/ops/add/binary_add_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.h>
 
 #include <webgpu/webgpu.h>
 
@@ -42,6 +43,11 @@ void add_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     alpha = static_cast<float>(graph.get_double(alpha_id));
   }
 
+  // Fold into an immediately preceding rms_norm dispatch when guards hold.
+  if (fusion::try_fuse_add(graph, in1_id, in2_id, alpha, out_id)) {
+    return;
+  }
+
   const auto& in1_tensor = graph.get_tensor(in1_id);
   const auto& in2_tensor = graph.get_tensor(in2_id);
   const auto& out_tensor = graph.get_tensor(out_id);
diff --git a/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp b/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp
index f593334cb2a..4c2267b0eb0 100644
--- a/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp
+++ b/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp
@@ -14,6 +14,7 @@
 #include <webgpu/webgpu.h>
 
 #include <cstdint>
+#include <memory>
 #include <stdexcept>
 #include <vector>
 
@@ -31,15 +32,95 @@ static_assert(
     sizeof(ChooseQParamsParams) == 16,
     "ChooseQParamsParams must match the WGSL Params struct (16 bytes)");
 
+struct ChooseQParamsState {
+  ChooseQParamsParams params;
+  std::vector<int64_t> output_dims;
+  utils::WgCount grid;
+};
+
+ChooseQParamsState make_choose_qparams_state(
+    WGPUDevice device,
+    const std::vector<int64_t>& input_dims,
+    uint32_t max_rows,
+    uint32_t reduce_size,
+    int32_t quant_min,
+    int32_t quant_max) {
+  if (input_dims.empty() || input_dims.back() != reduce_size) {
+    throw std::runtime_error(
+        "choose_qparams_affine: live reduce size mismatch");
+  }
+  const uint64_t numel = utils::numel_of(input_dims);
+  if (numel == 0u || numel % reduce_size != 0u) {
+    throw std::runtime_error("choose_qparams_affine: invalid live input numel");
+  }
+  const uint64_t rows = numel / reduce_size;
+  if (rows == 0u || rows > max_rows || rows > UINT32_MAX) {
+    throw std::runtime_error(
+        "choose_qparams_affine: live rows exceed the build-time max");
+  }
+
+  ChooseQParamsState state = {};
+  state.params = {
+      static_cast<uint32_t>(rows), reduce_size, quant_min, quant_max};
+  state.output_dims = input_dims;
+  state.output_dims.pop_back();
+  state.grid = utils::compute_2d_workgroup_count(
+      device,
+      utils::div_up(static_cast<uint32_t>(rows), 4u),
+      1u,
+      "choose_qparams_affine");
+  return state;
+}
+
 // torchao.choose_qparams_affine args (mirrors Vulkan ChooseQParams.cpp:158):
 // [input, mapping_type, block_size, target_dtype, quant_min, quant_max, eps,
-//  scale_dtype, zero_point_dtype, out_tuple(scale, zp)]. Routes to the per-row
-// (last-dim) path: one workgroup per row computes asymmetric scale/zp.
+//  scale_dtype, zero_point_dtype, keepdim, out_tuple(scale, zp)].
+// Routes to the per-row (last-dim) path.
 void choose_qparams_affine_impl(
     WebGPUGraph& graph,
     const std::vector<int>& args) {
+  if (args.size() != 11u) {
+    throw std::runtime_error(
+        "choose_qparams_affine: expected 10 inputs plus output");
+  }
   const int in_id = args.at(0);
-  const int out_list_id = args.at(args.size() - 1);
+  const int mapping_type_id = args.at(1);
+  const int block_size_id = args.at(2);
+  const int target_dtype_id = args.at(3);
+  const int quant_min_id = args.at(4);
+  const int quant_max_id = args.at(5);
+  const int eps_id = args.at(6);
+  const int scale_dtype_id = args.at(7);
+  const int zero_point_dtype_id = args.at(8);
+  const int keepdim_id = args.at(9);
+  const int out_list_id = args.at(10);
+
+  using VT = WebGPUGraph::ValueType;
+  if (graph.get_value_type(mapping_type_id) != VT::String ||
+      graph.get_string(mapping_type_id) != "ASYMMETRIC") {
+    throw std::runtime_error(
+        "choose_qparams_affine: only ASYMMETRIC mapping is supported");
+  }
+  if (graph.get_value_type(block_size_id) != VT::IntList ||
+      graph.get_value_type(quant_min_id) != VT::Int ||
+      graph.get_value_type(quant_max_id) != VT::Int ||
+      graph.get_value_type(keepdim_id) != VT::Bool ||
+      graph.get_value_type(out_list_id) != VT::ValueList) {
+    throw std::runtime_error("choose_qparams_affine: malformed scalar args");
+  }
+  // The current Vulkan serializer encodes torch.dtype values as Null and
+  // materializes the schema-default keepdim=false. The exact supported dtypes
+  // are therefore validated from the output tensors below.
+  if (graph.get_value_type(target_dtype_id) != VT::Null ||
+      graph.get_value_type(eps_id) != VT::Null ||
+      graph.get_value_type(scale_dtype_id) != VT::Null ||
+      graph.get_value_type(zero_point_dtype_id) != VT::Null) {
+    throw std::runtime_error(
+        "choose_qparams_affine: unsupported serialized dtype arguments");
+  }
+  if (graph.get_bool(keepdim_id)) {
+    throw std::runtime_error("choose_qparams_affine: keepdim must be false");
+  }
 
   const std::vector<int>& out_ids = graph.get_value_list(out_list_id);
   if (out_ids.size() != 2) {
@@ -63,50 +144,44 @@ void choose_qparams_affine_impl(
       zp_t.buffer == nullptr) {
     throw std::runtime_error("choose_qparams_affine: null buffer binding");
   }
-  if (in.dims.empty()) {
-    throw std::runtime_error("choose_qparams_affine: input has no dims");
+  if (in.dims.empty() || in.is_int || in.elem_size != sizeof(float) ||
+      scale_t.is_int || scale_t.elem_size != sizeof(float)) {
+    throw std::runtime_error(
+        "choose_qparams_affine: input and scale must be fp32");
   }
 
-  const uint64_t reduce_size = static_cast<uint64_t>(in.dims.back());
-  if (reduce_size == 0) {
-    throw std::runtime_error("choose_qparams_affine: last dim == 0");
+  if (in.dims.back() <= 0 ||
+      static_cast<uint64_t>(in.dims.back()) > UINT32_MAX) {
+    throw std::runtime_error("choose_qparams_affine: invalid last dimension");
   }
-  uint64_t in_numel = 1;
-  for (int64_t d : in.dims) {
-    in_numel *= static_cast<uint64_t>(d);
+  const uint64_t reduce_size = static_cast<uint64_t>(in.dims.back());
+  const uint64_t in_numel = utils::numel_of(in.dims);
+  if (in_numel > SIZE_MAX / sizeof(float)) {
+    throw std::runtime_error("choose_qparams_affine: input size overflows");
   }
   const uint64_t num_rows = in_numel / reduce_size;
-  if (num_rows == 0 || num_rows > UINT32_MAX || reduce_size > UINT32_MAX) {
+  if (in_numel % reduce_size != 0u || num_rows == 0 || num_rows > UINT32_MAX) {
     throw std::runtime_error("choose_qparams_affine: bad row/reduce shape");
   }
-  if (in.nbytes != in_numel * sizeof(float)) {
+  if (in.nbytes != static_cast<size_t>(in_numel) * sizeof(float)) {
     throw std::runtime_error("choose_qparams_affine: input must be fp32");
   }
   // scale is fp32[num_rows]; zp is int8[num_rows] (bound as array<u32>).
-  if (scale_t.nbytes != num_rows * sizeof(float)) {
+  if (scale_t.dims != zp_t.dims || utils::numel_of(scale_t.dims) != num_rows ||
+      scale_t.nbytes != num_rows * sizeof(float)) {
     throw std::runtime_error("choose_qparams_affine: scale must be fp32[rows]");
   }
-  // zp is int8[rows] (elem_size 1), packed 4-per-u32 in the shader. int8
-  // buffers are allocated max(nbytes, 4); M<=4 pads to one word, M%4==0 is
-  // word-exact. Other M (5,6,7,...) would overflow the M-byte buffer -> reject.
-  if (!zp_t.is_int8 || zp_t.nbytes != num_rows) {
+  // zp is int8[rows] (elem_size 1), packed 4-per-u32 in the shader. Buffers are
+  // allocated max(align4(nbytes), 4), so a ragged tail block's whole-word store
+  // lands in the pad, in-bounds; the shader clamps its row loop to num_rows.
+  if (!zp_t.is_int8 || zp_t.elem_size != sizeof(int8_t) ||
+      zp_t.nbytes != num_rows) {
     throw std::runtime_error("choose_qparams_affine: zp must be int8[rows]");
   }
-  if (num_rows > 4 && num_rows % 4 != 0) {
-    throw std::runtime_error(
-        "choose_qparams_affine: num_rows must be <=4 or a multiple of 4");
-  }
 
   // The kernel implements only the asymmetric, per-row (last-dim), int8 path;
   // validate the schema args it assumes and fail loud rather than silently
   // ignoring them (mirrors Vulkan, which consumes block_size + quant_min/max).
-  const int quant_min_id = args.at(4);
-  const int quant_max_id = args.at(5);
-  if (graph.get_value_type(quant_min_id) != WebGPUGraph::ValueType::Int ||
-      graph.get_value_type(quant_max_id) != WebGPUGraph::ValueType::Int) {
-    throw std::runtime_error(
-        "choose_qparams_affine: quant_min/quant_max must be int scalars");
-  }
   const int64_t quant_min = graph.get_int(quant_min_id);
   const int64_t quant_max = graph.get_int(quant_max_id);
   if (quant_min != -128 || quant_max != 127) {
@@ -114,13 +189,8 @@ void choose_qparams_affine_impl(
         "choose_qparams_affine: only the int8 range [-128, 127] is supported");
   }
   // Per-row fast path: block_size must be [1, ..., 1, reduce_size].
-  const int block_size_id = args.at(2);
-  if (graph.get_value_type(block_size_id) != WebGPUGraph::ValueType::IntList) {
-    throw std::runtime_error(
-        "choose_qparams_affine: block_size must be an int list");
-  }
   const std::vector<int64_t>& block_size = graph.get_int_list(block_size_id);
-  if (block_size.empty() ||
+  if (block_size.size() != in.dims.size() ||
       block_size.back() != static_cast<int64_t>(reduce_size)) {
     throw std::runtime_error(
         "choose_qparams_affine: block_size must reduce the last dim");
@@ -132,26 +202,25 @@ void choose_qparams_affine_impl(
     }
   }
 
-  ChooseQParamsParams params = {};
-  params.num_rows = static_cast<uint32_t>(num_rows);
-  params.reduce_size = static_cast<uint32_t>(reduce_size);
-  params.quant_min = static_cast<int32_t>(quant_min);
-  params.quant_max = static_cast<int32_t>(quant_max);
+  const uint32_t max_rows = static_cast<uint32_t>(num_rows);
+  const uint32_t reduce_size_u32 = static_cast<uint32_t>(reduce_size);
+  const ChooseQParamsState initial_state = make_choose_qparams_state(
+      device,
+      in.dims,
+      max_rows,
+      reduce_size_u32,
+      static_cast<int32_t>(quant_min),
+      static_cast<int32_t>(quant_max));
 
   uint32_t wg_size =
       utils::clamp_workgroup_size(device, kChooseQparamsAffineWorkgroupSizeX);
-  // One workgroup per block of 4 rows (wg_size threads cooperate per row); the
-  // block packs its 4 int8 zps into one u32. 2D-fold lifts the 65535 grid cap.
-  const uint32_t num_blocks = static_cast<uint32_t>((num_rows + 3) / 4);
-  utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
-      device, num_blocks, 1, "choose_qparams_affine");
 
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
   wg_size_constant.value = static_cast<double>(wg_size);
 
-  WGPUBuffer params_buf =
-      utils::make_uniform(device, &params, sizeof(ChooseQParamsParams));
+  WGPUBuffer params_buf = utils::make_uniform(
+      device, &initial_state.params, sizeof(ChooseQParamsParams));
   graph.add_uniform_buffer_bytes(sizeof(ChooseQParamsParams));
 
   utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
@@ -174,12 +243,58 @@ void choose_qparams_affine_impl(
       &wg_size_constant,
       1);
 
-  graph.add_dispatch(
+  const size_t dispatch_index = graph.add_dispatch(
       {bundle.pipeline,
        bundle.bind_group,
-       workgroup_count.x,
+       initial_state.grid.x,
        "choose_qparams_affine",
-       workgroup_count.y});
+       initial_state.grid.y});
+
+  auto producer_elided = std::make_shared<bool>(false);
+  graph.add_tensor_resize_hook(
+      in_id,
+      [in_id,
+       scale_id,
+       zp_id,
+       max_rows,
+       reduce_size_u32,
+       quant_min = static_cast<int32_t>(quant_min),
+       quant_max = static_cast<int32_t>(quant_max),
+       dispatch_index,
+       producer_elided,
+       params_buf](WebGPUGraph& g) {
+        const ChooseQParamsState state = make_choose_qparams_state(
+            g.device(),
+            g.cur_dims(in_id),
+            max_rows,
+            reduce_size_u32,
+            quant_min,
+            quant_max);
+        wgpuQueueWriteBuffer(
+            g.queue(), params_buf, 0, &state.params, sizeof(state.params));
+        auto& dispatch = g.dispatch_at(dispatch_index);
+        dispatch.workgroup_count_x =
+            utils::cqp_resize_workgroups(*producer_elided, state.grid.x);
+        dispatch.workgroup_count_y =
+            utils::cqp_resize_workgroups(*producer_elided, state.grid.y);
+        g.set_cur_dims(scale_id, state.output_dims);
+        g.set_cur_dims(zp_id, state.output_dims);
+      });
+
+  WebGPUGraph::CqpFusionSite site = {};
+  site.input_id = in_id;
+  site.scales_id = scale_id;
+  site.zero_points_id = zp_id;
+  site.rows = max_rows;
+  site.row_width = reduce_size_u32;
+  site.quant_min = quant_min;
+  site.quant_max = quant_max;
+  site.dispatch_index = dispatch_index;
+  site.input_buffer = in.buffer;
+  site.scales_buffer = scale_t.buffer;
+  site.zero_points_buffer = zp_t.buffer;
+  site.producer_elided = std::move(producer_elided);
+  graph.offer_cqp_fusion_site(std::move(site));
 
   graph.own_uniform_buffer(params_buf);
 }
diff --git a/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl
index ff0c9002ada..d3598fb0ec6 100644
--- a/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl
+++ b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl
@@ -21,6 +21,12 @@ var<workgroup> part_max: array<f32, 256>;
 
 const SMALL_SCALE_THRESHOLD: f32 = 6.1e-5;
 
+fn reciprocal_is_infinite(value: f32) -> bool {
+  // WGSL has no portable isinf builtin. This exponent/mantissa check is the
+  // exact f32 equivalent used for Vulkan's isinf(1.0 / scale) condition.
+  return (bitcast<u32>(1.0 / value) & 0x7fffffffu) == 0x7f800000u;
+}
+
 @compute @workgroup_size(wg_size)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
@@ -68,7 +74,7 @@ fn main(
       mn = min(mn, 0.0);
       mx = max(mx, 0.0);
       var scale = (mx - mn) / (qmax - qmin);
-      if (scale == 0.0) {
+      if (scale == 0.0 || reciprocal_is_infinite(scale)) {
         scale = 0.1;
       }
       if (scale < SMALL_SCALE_THRESHOLD) {
diff --git a/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h
index bf2b027b906..4f42e4a027e 100644
--- a/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h
+++ b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from choose_qparams_affine.wgsl - DO NOT EDIT.
-// wgsl-sha256: 45b55ec7c432d5fbd9a7c9716b569c9c654e7ea566b6530c4b78f2924daa116f
+// wgsl-sha256: 15c9a27671ac6a4cfec24fb75389bd7aa2126dde1511b36bdd678f3338ba71e8
 inline constexpr const char* kChooseQparamsAffineWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> scales_out: array<f32>;
@@ -38,6 +38,12 @@ var<workgroup> part_max: array<f32, 256>;
 
 const SMALL_SCALE_THRESHOLD: f32 = 6.1e-5;
 
+fn reciprocal_is_infinite(value: f32) -> bool {
+  // WGSL has no portable isinf builtin. This exponent/mantissa check is the
+  // exact f32 equivalent used for Vulkan's isinf(1.0 / scale) condition.
+  return (bitcast<u32>(1.0 / value) & 0x7fffffffu) == 0x7f800000u;
+}
+
 @compute @workgroup_size(wg_size)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
@@ -85,7 +91,7 @@ fn main(
       mn = min(mn, 0.0);
       mx = max(mx, 0.0);
       var scale = (mx - mn) / (qmax - qmin);
-      if (scale == 0.0) {
+      if (scale == 0.0 || reciprocal_is_infinite(scale)) {
         scale = 0.1;
       }
       if (scale < SMALL_SCALE_THRESHOLD) {
diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp
index b5ea0cdca4c..7bf5ee83b14 100644
--- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp
+++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp
@@ -13,16 +13,20 @@
 
 #include <webgpu/webgpu.h>
 
+#include <algorithm>
 #include <cstdint>
 #include <cstring>
+#include <limits>
+#include <numeric>
 #include <stdexcept>
+#include <utility>
 #include <vector>
 
 namespace executorch::backends::webgpu {
 
 namespace {
 
-// Uniform layout matching the WGSL Params struct (16-byte aligned, 32 bytes).
+// Uniform layout matching the WGSL Params struct (16-byte aligned, 48 bytes).
 struct EmbeddingParams {
   uint32_t embed_dim;
   uint32_t blocks_per_row;
@@ -32,10 +36,14 @@ struct EmbeddingParams {
   uint32_t bytes_per_row;
   uint32_t total_blocks;
   uint32_t is_linear_weight;
+  uint32_t row_lo;
+  uint32_t rows_in_chunk;
+  uint32_t pad0;
+  uint32_t pad1;
 };
 static_assert(
-    sizeof(EmbeddingParams) == 32,
-    "EmbeddingParams must be 32 bytes");
+    sizeof(EmbeddingParams) == 48,
+    "EmbeddingParams must be 48 bytes");
 
 struct EmbeddingLayout {
   uint32_t embed_dim;
@@ -46,10 +54,27 @@ struct EmbeddingLayout {
   bool is_linear_weight;
 };
 
+struct EmbeddingChunkSpec {
+  uint32_t row_lo;
+  uint32_t rows;
+  uint64_t weight_offset;
+  uint64_t weight_size;
+  uint64_t scales_offset;
+  uint64_t scales_size;
+};
+
+struct EmbeddingChunkRuntime {
+  EmbeddingChunkSpec spec;
+  size_t dispatch_index;
+  WGPUBuffer params_buffer;
+};
+
 EmbeddingParams make_embedding_params(
     const EmbeddingLayout& layout,
     uint32_t num_indices,
-    uint32_t total_blocks) {
+    uint32_t total_blocks,
+    uint32_t row_lo,
+    uint32_t rows_in_chunk) {
   return {
       layout.embed_dim,
       layout.blocks_per_row,
@@ -58,7 +83,89 @@ EmbeddingParams make_embedding_params(
       layout.groups_per_row,
       layout.bytes_per_row,
       total_blocks,
-      layout.is_linear_weight ? 1u : 0u};
+      layout.is_linear_weight ? 1u : 0u,
+      row_lo,
+      rows_in_chunk,
+      0u,
+      0u};
+}
+
+uint64_t checked_lcm(uint64_t a, uint64_t b) {
+  const uint64_t divisor = std::gcd(a, b);
+  if (a > std::numeric_limits<uint64_t>::max() / (b / divisor)) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: row-alignment quantum overflows");
+  }
+  return a * (b / divisor);
+}
+
+std::vector<EmbeddingChunkSpec> make_embedding_chunks(
+    uint64_t max_binding_bytes,
+    uint64_t min_offset_alignment,
+    uint64_t max_buffer_bytes,
+    uint32_t vocab_rows,
+    uint64_t weight_bytes_per_row,
+    uint64_t scales_bytes_per_row,
+    uint64_t weight_buffer_bytes,
+    uint64_t scales_buffer_bytes) {
+  if (max_binding_bytes == 0u || min_offset_alignment == 0u ||
+      max_buffer_bytes == 0u || vocab_rows == 0u ||
+      weight_bytes_per_row == 0u || scales_bytes_per_row == 0u) {
+    throw std::runtime_error("WebGPU embedding_q4gsw: invalid chunking limits");
+  }
+  if (weight_buffer_bytes > max_buffer_bytes ||
+      scales_buffer_bytes > max_buffer_bytes) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: tensor exceeds maxBufferSize");
+  }
+  if (weight_buffer_bytes <= max_binding_bytes &&
+      scales_buffer_bytes <= max_binding_bytes) {
+    return {{0u, vocab_rows, 0u, weight_buffer_bytes, 0u, scales_buffer_bytes}};
+  }
+
+  const uint64_t weight_row_quantum = min_offset_alignment /
+      std::gcd(min_offset_alignment, weight_bytes_per_row);
+  const uint64_t scales_row_quantum = min_offset_alignment /
+      std::gcd(min_offset_alignment, scales_bytes_per_row);
+  const uint64_t row_quantum =
+      checked_lcm(weight_row_quantum, scales_row_quantum);
+  const uint64_t max_rows = std::min(
+      max_binding_bytes / weight_bytes_per_row,
+      max_binding_bytes / scales_bytes_per_row);
+  const uint64_t rows_per_chunk = (max_rows / row_quantum) * row_quantum;
+  if (rows_per_chunk == 0u) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: no aligned row chunk fits binding limit");
+  }
+
+  std::vector<EmbeddingChunkSpec> chunks;
+  for (uint64_t row_lo = 0u; row_lo < vocab_rows; row_lo += rows_per_chunk) {
+    const uint64_t rows =
+        std::min<uint64_t>(rows_per_chunk, vocab_rows - row_lo);
+    const uint64_t weight_offset = row_lo * weight_bytes_per_row;
+    const uint64_t scales_offset = row_lo * scales_bytes_per_row;
+    const uint64_t weight_size = rows * weight_bytes_per_row;
+    const uint64_t scales_size = rows * scales_bytes_per_row;
+    if (weight_offset % min_offset_alignment != 0u ||
+        scales_offset % min_offset_alignment != 0u ||
+        weight_size > max_binding_bytes || scales_size > max_binding_bytes ||
+        weight_offset > weight_buffer_bytes ||
+        weight_size > weight_buffer_bytes - weight_offset ||
+        scales_offset > scales_buffer_bytes ||
+        scales_size > scales_buffer_bytes - scales_offset ||
+        row_lo > UINT32_MAX || rows > UINT32_MAX) {
+      throw std::runtime_error(
+          "WebGPU embedding_q4gsw: invalid aligned chunk binding");
+    }
+    chunks.push_back(
+        {static_cast<uint32_t>(row_lo),
+         static_cast<uint32_t>(rows),
+         weight_offset,
+         weight_size,
+         scales_offset,
+         scales_size});
+  }
+  return chunks;
 }
 
 // Resize hook body: recompute counts/dispatch; out = indices dims +
@@ -69,8 +176,7 @@ void resize_embedding_q4gsw(
     int out_id,
     const EmbeddingLayout& layout,
     uint32_t wg_size,
-    size_t dispatch_idx,
-    WGPUBuffer params_buf) {
+    const std::vector<EmbeddingChunkRuntime>& chunks) {
   const auto& id = g.cur_dims(indices_id);
   const uint64_t ni = utils::numel_of(id);
   if (ni == 0) {
@@ -84,15 +190,23 @@ void resize_embedding_q4gsw(
   std::vector<int64_t> od = id;
   od.push_back(static_cast<int64_t>(layout.embed_dim));
   g.set_cur_dims(out_id, od);
-  EmbeddingParams p = make_embedding_params(
-      layout, static_cast<uint32_t>(ni), static_cast<uint32_t>(total_blocks));
-  wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
-  g.dispatch_at(dispatch_idx).workgroup_count_x =
-      utils::compute_1d_workgroup_count(
-          g.device(),
-          static_cast<uint32_t>(total_blocks),
-          wg_size,
-          "embedding_q4gsw(resize)");
+  const utils::WgCount grid = utils::compute_2d_workgroup_count(
+      g.device(),
+      static_cast<uint32_t>(total_blocks),
+      wg_size,
+      "embedding_q4gsw(resize)");
+  for (const EmbeddingChunkRuntime& chunk : chunks) {
+    const EmbeddingParams p = make_embedding_params(
+        layout,
+        static_cast<uint32_t>(ni),
+        static_cast<uint32_t>(total_blocks),
+        chunk.spec.row_lo,
+        chunk.spec.rows);
+    wgpuQueueWriteBuffer(g.queue(), chunk.params_buffer, 0, &p, sizeof(p));
+    WebGPUDispatch& dispatch = g.dispatch_at(chunk.dispatch_index);
+    dispatch.workgroup_count_x = grid.x;
+    dispatch.workgroup_count_y = grid.y;
+  }
 }
 
 // arg order mirrors Vulkan EmbeddingQ4gsw.cpp.
@@ -131,8 +245,18 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     throw std::runtime_error("WebGPU embedding_q4gsw: malformed dims");
   }
 
+  if (out.dims.back() <= 0 ||
+      static_cast<uint64_t>(out.dims.back()) > UINT32_MAX ||
+      weight.dims[0] <= 0 ||
+      static_cast<uint64_t>(weight.dims[0]) > UINT32_MAX ||
+      weight.dims[1] <= 0 || scales.dims[0] <= 0 || scales.dims[1] <= 0 ||
+      static_cast<uint64_t>(scales.dims[1]) > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: dimensions exceed supported range");
+  }
   const uint32_t embed_dim = static_cast<uint32_t>(out.dims.back());
-  if (embed_dim == 0 || embed_dim % 32 != 0) {
+  const uint32_t vocab_rows = static_cast<uint32_t>(weight.dims[0]);
+  if (embed_dim % 32u != 0u) {
     throw std::runtime_error(
         "WebGPU embedding_q4gsw: embed_dim must be a nonzero multiple of 32");
   }
@@ -145,12 +269,17 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   if (graph.get_value_type(group_size_id) == WebGPUGraph::ValueType::Int) {
     group_size = graph.get_int(group_size_id);
   }
-  if (group_size <= 0) {
-    throw std::runtime_error("WebGPU embedding_q4gsw: group_size <= 0");
+  if (group_size <= 0 || static_cast<uint64_t>(group_size) > UINT32_MAX) {
+    throw std::runtime_error("WebGPU embedding_q4gsw: group_size out of range");
   }
 
   // Leading index dims flatten row-major (mirrors Vulkan num_indices).
   const uint64_t out_numel = utils::numel_of(out.dims);
+  if (out_numel % embed_dim != 0u || out_numel / embed_dim == 0u ||
+      out_numel / embed_dim > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: invalid number of indices");
+  }
   const uint32_t num_indices = static_cast<uint32_t>(out_numel / embed_dim);
   const uint32_t groups_per_row = static_cast<uint32_t>(scales.dims[1]);
   const uint32_t blocks_per_row = embed_dim / 32u;
@@ -161,6 +290,10 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     throw std::runtime_error(
         "WebGPU embedding_q4gsw: groups_per_row * group_size != embed_dim");
   }
+  if (scales.dims[0] != weight.dims[0]) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: weight/scales vocab rows differ");
+  }
   if (weight.buffer == nullptr || scales.buffer == nullptr ||
       indices.buffer == nullptr || out.buffer == nullptr) {
     throw std::runtime_error("WebGPU embedding_q4gsw: null buffer binding");
@@ -170,11 +303,13 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const uint64_t indices_numel = utils::numel_of(indices.dims);
   const uint64_t weight_numel = utils::numel_of(weight.dims);
   const uint64_t scales_numel = utils::numel_of(scales.dims);
-  if (indices_numel != num_indices ||
+  if (indices_numel != num_indices || !indices.is_int ||
+      indices.elem_size != sizeof(int32_t) ||
       indices.nbytes != indices_numel * sizeof(int32_t) ||
       weight.nbytes != weight_numel ||
-      scales.nbytes != scales_numel * sizeof(float) ||
-      out.nbytes != out_numel * sizeof(float)) {
+      weight_numel != static_cast<uint64_t>(vocab_rows) * bytes_per_row ||
+      scales_numel != static_cast<uint64_t>(vocab_rows) * groups_per_row ||
+      !utils::is_fp32_tensor(scales) || !utils::is_fp32_tensor(out)) {
     throw std::runtime_error(
         "WebGPU embedding_q4gsw: dtype/byte-size mismatch "
         "(indices int32, weight uint8, scales/out fp32)");
@@ -184,10 +319,41 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         "WebGPU embedding_q4gsw: total_blocks exceeds uint32 dispatch range");
   }
 
+  std::vector<int64_t> expected_out_dims = indices.dims;
+  expected_out_dims.push_back(embed_dim);
+  if (out.dims != expected_out_dims) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: output shape must be indices + embed_dim");
+  }
+
+  WGPULimits limits = {};
+  if (wgpuDeviceGetLimits(device, &limits) != WGPUStatus_Success ||
+      limits.maxStorageBufferBindingSize == 0u || limits.maxBufferSize == 0u ||
+      limits.minStorageBufferOffsetAlignment == 0u) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: failed to query binding limits");
+  }
+  if (indices.nbytes > limits.maxStorageBufferBindingSize ||
+      out.nbytes > limits.maxStorageBufferBindingSize ||
+      indices.nbytes > limits.maxBufferSize ||
+      out.nbytes > limits.maxBufferSize) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: indices/output exceed binding limits");
+  }
+  const std::vector<EmbeddingChunkSpec> chunk_specs = make_embedding_chunks(
+      limits.maxStorageBufferBindingSize,
+      limits.minStorageBufferOffsetAlignment,
+      limits.maxBufferSize,
+      vocab_rows,
+      bytes_per_row,
+      static_cast<uint64_t>(groups_per_row) * sizeof(float),
+      weight.nbytes,
+      scales.nbytes);
+
   // 1D dispatch: one thread per 32-dim block; validate before any alloc.
   const uint32_t wg_size =
       utils::clamp_workgroup_size(device, kEmbeddingQ4gswWorkgroupSizeX);
-  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+  const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
       device, static_cast<uint32_t>(total_blocks), wg_size, "embedding_q4gsw");
 
   const EmbeddingLayout layout = {
@@ -197,64 +363,64 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       groups_per_row,
       bytes_per_row,
       is_linear};
-  EmbeddingParams params = make_embedding_params(
-      layout, num_indices, static_cast<uint32_t>(total_blocks));
-
-  WGPUBufferDescriptor uniform_desc = {};
-  uniform_desc.size = sizeof(EmbeddingParams);
-  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
-  uniform_desc.mappedAtCreation = true;
-  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);
-  void* mapped =
-      wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(EmbeddingParams));
-  std::memcpy(mapped, &params, sizeof(EmbeddingParams));
-  wgpuBufferUnmap(uniform_buffer);
-  graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams));
-
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
   wg_size_constant.value = static_cast<double>(wg_size);
 
-  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
-      device,
-      kEmbeddingQ4gswWGSL,
-      {
-          {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
-          {1,
-           WGPUBufferBindingType_ReadOnlyStorage,
-           indices.buffer,
-           indices.nbytes},
-          {2,
-           WGPUBufferBindingType_ReadOnlyStorage,
-           weight.buffer,
-           weight.nbytes},
-          {3,
-           WGPUBufferBindingType_ReadOnlyStorage,
-           scales.buffer,
-           scales.nbytes},
-          {4,
-           WGPUBufferBindingType_Uniform,
-           uniform_buffer,
-           sizeof(EmbeddingParams)},
-      },
-      &wg_size_constant,
-      1);
-
-  const size_t dispatch_idx = graph.add_dispatch(
-      {bundle.pipeline, bundle.bind_group, workgroup_count, "embedding_q4gsw"});
+  std::vector<EmbeddingChunkRuntime> chunks;
+  chunks.reserve(chunk_specs.size());
+  for (const EmbeddingChunkSpec& chunk : chunk_specs) {
+    const EmbeddingParams params = make_embedding_params(
+        layout,
+        num_indices,
+        static_cast<uint32_t>(total_blocks),
+        chunk.row_lo,
+        chunk.rows);
+    WGPUBuffer params_buffer = graph.create_params_buffer(params);
+    graph.add_uniform_buffer_bytes(sizeof(params));
+
+    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+        device,
+        kEmbeddingQ4gswWGSL,
+        {
+            {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
+            {1,
+             WGPUBufferBindingType_ReadOnlyStorage,
+             indices.buffer,
+             indices.nbytes},
+            {2,
+             WGPUBufferBindingType_ReadOnlyStorage,
+             weight.buffer,
+             chunk.weight_size,
+             chunk.weight_offset},
+            {3,
+             WGPUBufferBindingType_ReadOnlyStorage,
+             scales.buffer,
+             chunk.scales_size,
+             chunk.scales_offset},
+            {4,
+             WGPUBufferBindingType_Uniform,
+             params_buffer,
+             sizeof(EmbeddingParams)},
+        },
+        &wg_size_constant,
+        1);
+
+    const size_t dispatch_index = graph.add_dispatch(
+        {bundle.pipeline,
+         bundle.bind_group,
+         workgroup_count.x,
+         "embedding_q4gsw",
+         workgroup_count.y});
+    chunks.push_back({chunk, dispatch_index, params_buffer});
+  }
 
   // Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim].
-  WGPUBuffer params_buf = uniform_buffer;
   graph.add_tensor_resize_hook(
       indices_id,
-      [indices_id, out_id, layout, wg_size, dispatch_idx, params_buf](
-          WebGPUGraph& g) {
-        resize_embedding_q4gsw(
-            g, indices_id, out_id, layout, wg_size, dispatch_idx, params_buf);
+      [indices_id, out_id, layout, wg_size, chunks](WebGPUGraph& g) {
+        resize_embedding_q4gsw(g, indices_id, out_id, layout, wg_size, chunks);
       });
-
-  // Graph owns it so the resize hook can rewrite it; freed in the dtor.
-  graph.own_uniform_buffer(uniform_buffer);
 }
 
 } // namespace
diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw.wgsl b/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw.wgsl
index fecb5f1e28a..b0ad579c3ad 100644
--- a/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw.wgsl
+++ b/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw.wgsl
@@ -12,6 +12,13 @@ struct Params {
   bytes_per_row: u32,
   total_blocks: u32,
   is_linear_weight: u32,
+  // This dispatch owns vocab rows [row_lo, row_lo + rows_in_chunk). t_weight
+  // and t_scales are bound at that chunk's byte offsets, so both use row-local
+  // indices.
+  row_lo: u32,
+  rows_in_chunk: u32,
+  pad0: u32,
+  pad1: u32,
 }
 @group(0) @binding(4) var<uniform> params: Params;
 
@@ -19,8 +26,11 @@ override wg_size: u32 = 64u;
 
 // One thread per 32-dim block of one gathered row (flat-buffer weight path).
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let block = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  // 2D-spill: total_blocks can exceed the 65535 per-dim grid cap.
+  let block = gid.x + gid.y * (num_workgroups.x * wg_size);
   if (block >= params.total_blocks) {
     return;
   }
@@ -29,7 +39,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
   // token assumed in-range (mirrors Vulkan; no vocab clamp).
   let token = u32(t_indices[indices_idx]);
-  let row_byte_base = token * params.bytes_per_row;
+  // A vocab-sized 4-bit table can exceed maxStorageBufferBindingSize, so it is
+  // bound one chunk of rows per dispatch. Each token belongs to exactly one
+  // chunk, so every output row is still written exactly once.
+  if (token < params.row_lo || token - params.row_lo >= params.rows_in_chunk) {
+    return;
+  }
+  let row_byte_base = (token - params.row_lo) * params.bytes_per_row;
   let out_base = indices_idx * params.embed_dim + base_dim;
 
   for (var t: u32 = 0u; t < 32u; t = t + 1u) {
@@ -46,7 +62,9 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
       nib = b & 0x0Fu;          // low nibble
     }
     let q = f32(i32(nib) - 8); // +8-shifted on pack; recover signed [-8,7]
-    let scale = t_scales[token * params.groups_per_row + dim / params.group_size];
+    let scale = t_scales[
+        (token - params.row_lo) * params.groups_per_row +
+        dim / params.group_size];
     t_out[out_base + t] = q * scale;
   }
 }
diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw_wgsl.h b/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw_wgsl.h
index db26795a021..775944a4371 100644
--- a/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw_wgsl.h
+++ b/backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from embedding_q4gsw.wgsl - DO NOT EDIT.
-// wgsl-sha256: 94da1061b49b62556a79020182a4989439a7c51f919e83d577536c5b6d25f487
+// wgsl-sha256: 6a7c585b3e332a0c21604762d0687744365a7a3d94b15fa00ae4c0231749e956
 inline constexpr const char* kEmbeddingQ4gswWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_indices: array<i32>;
@@ -29,6 +29,13 @@ struct Params {
   bytes_per_row: u32,
   total_blocks: u32,
   is_linear_weight: u32,
+  // This dispatch owns vocab rows [row_lo, row_lo + rows_in_chunk). t_weight
+  // and t_scales are bound at that chunk's byte offsets, so both use row-local
+  // indices.
+  row_lo: u32,
+  rows_in_chunk: u32,
+  pad0: u32,
+  pad1: u32,
 }
 @group(0) @binding(4) var<uniform> params: Params;
 
@@ -36,8 +43,11 @@ override wg_size: u32 = 64u;
 
 // One thread per 32-dim block of one gathered row (flat-buffer weight path).
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let block = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  // 2D-spill: total_blocks can exceed the 65535 per-dim grid cap.
+  let block = gid.x + gid.y * (num_workgroups.x * wg_size);
   if (block >= params.total_blocks) {
     return;
   }
@@ -46,7 +56,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
   // token assumed in-range (mirrors Vulkan; no vocab clamp).
   let token = u32(t_indices[indices_idx]);
-  let row_byte_base = token * params.bytes_per_row;
+  // A vocab-sized 4-bit table can exceed maxStorageBufferBindingSize, so it is
+  // bound one chunk of rows per dispatch. Each token belongs to exactly one
+  // chunk, so every output row is still written exactly once.
+  if (token < params.row_lo || token - params.row_lo >= params.rows_in_chunk) {
+    return;
+  }
+  let row_byte_base = (token - params.row_lo) * params.bytes_per_row;
   let out_base = indices_idx * params.embed_dim + base_dim;
 
   for (var t: u32 = 0u; t < 32u; t = t + 1u) {
@@ -63,7 +79,9 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
       nib = b & 0x0Fu;          // low nibble
     }
     let q = f32(i32(nib) - 8); // +8-shifted on pack; recover signed [-8,7]
-    let scale = t_scales[token * params.groups_per_row + dim / params.group_size];
+    let scale = t_scales[
+        (token - params.row_lo) * params.groups_per_row +
+        dim / params.group_size];
     t_out[out_base + t] = q * scale;
   }
 }
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp
index 92a53cb7574..fff154e4ba1 100644
--- a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp
@@ -10,46 +10,74 @@
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h>
 
 #include <webgpu/webgpu.h>
 
+#include <algorithm>
 #include <cmath>
 #include <cstdint>
+#include <limits>
 #include <stdexcept>
+#include <string>
+#include <vector>
 
 namespace executorch::backends::webgpu {
 
 namespace {
 
+enum class SdpaLayout : uint32_t {
+  BHSD = 0,
+  BSHD = 1,
+};
+
+enum class MaskMode : uint32_t {
+  None = 0,
+  Rank2 = 1,
+  Expanded = 2,
+};
+
 struct QkParams {
   uint32_t B;
-  uint32_t H;
+  uint32_t Hq;
+  uint32_t Hkv;
   uint32_t S_q;
   uint32_t S_kv;
   uint32_t D;
+  uint32_t g;
   uint32_t has_mask;
+  uint32_t mask_mode;
+  uint32_t layout;
   uint32_t _pad0;
   float scale;
+  // This dispatch owns batch-head pairs [bh_lo, bh_lo + bh_count). The scratch
+  // is indexed relative to bh_lo; q/k/v stay absolute.
+  uint32_t bh_lo;
+  uint32_t bh_count;
+  uint32_t elide_masked_qk;
+  uint32_t _pad2;
 };
-static_assert(sizeof(QkParams) == 32, "QkParams must be 32 bytes");
+static_assert(sizeof(QkParams) == 64, "QkParams must be 64 bytes");
 
 struct AvParams {
   uint32_t B;
-  uint32_t H;
+  uint32_t Hq;
+  uint32_t Hkv;
   uint32_t S_q;
   uint32_t S_kv;
   uint32_t D;
+  uint32_t g;
+  uint32_t layout;
+  uint32_t bh_lo;
+  uint32_t bh_count;
   uint32_t _pad0;
   uint32_t _pad1;
-  uint32_t _pad2;
 };
-static_assert(sizeof(AvParams) == 32, "AvParams must be 32 bytes");
+static_assert(sizeof(AvParams) == 48, "AvParams must be 48 bytes");
 
-// Mirrors the Params struct in sdpa_softmax.wgsl (file-local in Sdpa.cpp, so
-// re-declared here for the reuse).
 struct SoftmaxParams {
   uint32_t num_rows;
   uint32_t row_width;
@@ -58,265 +86,784 @@ struct SoftmaxParams {
 };
 static_assert(sizeof(SoftmaxParams) == 16, "SoftmaxParams must be 16 bytes");
 
-// aten op: et_vk.sdpa.default. Args: [q, k, v, attn_mask, scale, out] (mirrors
-// Vulkan fused_sdpa_impl, SDPA.cpp). Non-causal, no KV-cache; all tensors are
-// DSHB [B, H, S, D], row-major. Three dispatches: QK (scaled, optional additive
-// mask) -> softmax (reused sdpa_softmax.wgsl) -> AV.
-void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector<int>& args) {
-  if (args.size() != 6) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: expected 6 args");
-  }
-  const int q_id = args.at(0);
-  const int k_id = args.at(1);
-  const int v_id = args.at(2);
-  const int mask_id = args.at(3);
-  const int scale_id = args.at(4);
-  const int out_id = args.at(5);
+struct TensorShape {
+  uint32_t B;
+  uint32_t H;
+  uint32_t S;
+  uint32_t D;
+};
 
-  WGPUDevice device = graph.device();
+struct SdpaShape {
+  uint32_t B;
+  uint32_t Hq;
+  uint32_t Hkv;
+  uint32_t S_q;
+  uint32_t S_kv;
+  uint32_t D;
+  uint32_t g;
+};
 
-  const auto& q = graph.get_tensor(q_id);
-  const auto& k = graph.get_tensor(k_id);
-  const auto& v = graph.get_tensor(v_id);
-  const auto& out = graph.get_tensor(out_id);
+// One recorded dispatch set owns a batch-head range, its uniforms, and either
+// one fixed QK dispatch or a dynamic row/entry route followed by softmax + AV.
+struct ChunkRecord {
+  uint32_t bh_lo;
+  uint32_t bh_count;
+  WGPUBuffer qk_params;
+  WGPUBuffer softmax_params;
+  WGPUBuffer av_params;
+  size_t qk_dispatch;
+  size_t qk_route_group;
+  bool fixed_qk_entry;
+  bool dual_qk;
+  size_t softmax_dispatch;
+  size_t av_dispatch;
+};
 
-  if (q.dims.size() < 3) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: q rank < 3");
-  }
-  const int rank = static_cast<int>(q.dims.size());
-  const uint32_t D = static_cast<uint32_t>(q.dims[rank - 1]);
-  const uint32_t S_q = static_cast<uint32_t>(q.dims[rank - 2]);
-  const uint32_t H = static_cast<uint32_t>(q.dims[rank - 3]);
-  if (D == 0 || S_q == 0 || H == 0) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: zero D/S_q/H");
-  }
-  // QK/AV kernels view q/k/v/out as vec4<f32> over D; every model in scope
-  // (Whisper/Voxtral/DaViT/BART/Hiera) uses D=64 or 128, always %4==0.
-  utils::check_vec4_aligned(D, "et_vk.sdpa", "D");
-  const uint64_t q_numel = utils::check_fp32(q, "et_vk.sdpa", "q");
-  const uint32_t B = static_cast<uint32_t>(q_numel / (uint64_t(H) * S_q * D));
-
-  // Asymmetric seq supported (S_q != S_kv, e.g. Hiera pooled query): k/v carry
-  // S_kv, q carries S_q; B/H/D must match across q/k/v. out is [B, H, S_q, D].
-  // When S_q == S_kv this is plain self-attention (bit-identical to before).
-  if (k.dims != v.dims) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: k/v shape mismatch");
-  }
-  if (k.dims.size() != q.dims.size()) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: q/k rank mismatch");
-  }
-  if (static_cast<uint32_t>(k.dims[rank - 1]) != D ||
-      static_cast<uint32_t>(k.dims[rank - 3]) != H) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: q/k/v must share H and D");
-  }
-  const uint32_t S_kv = static_cast<uint32_t>(k.dims[rank - 2]);
-  if (S_kv == 0) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: zero S_kv");
-  }
-  // Leading (batch) dims must agree across q/k/v.
-  for (int d = 0; d < rank - 3; ++d) {
-    if (k.dims[d] != q.dims[d]) {
-      throw std::runtime_error("WebGPU et_vk.sdpa: q/k batch dims mismatch");
-    }
+struct LiveState {
+  QkParams qk;
+  AvParams av;
+  SoftmaxParams softmax;
+  utils::WgCount qk_row_grid;
+  utils::WgCount qk_entry_grid;
+  utils::WgCount softmax_grid;
+  utils::WgCount av_grid;
+  bool use_qk_entry;
+};
+
+constexpr uint32_t kQkWorkgroupSize = 64;
+constexpr uint32_t kAvWorkgroupSize = 64;
+constexpr uint32_t kQkTileM = 8;
+constexpr uint32_t kQkTileN = 4;
+// One size serves both QK kernels; pin it so a .wgsl retune cannot drift.
+static_assert(
+    kQkWorkgroupSize == kEtVkSdpaQkWorkgroupSizeX &&
+        kQkWorkgroupSize == kEtVkSdpaQkEntryWorkgroupSizeX &&
+        kQkWorkgroupSize == kEtVkSdpaQkEntryExactWorkgroupSizeX,
+    "QK host workgroup size must match generated QK shader constants");
+static_assert(
+    kAvWorkgroupSize == kEtVkSdpaAvWorkgroupSizeX,
+    "AV host workgroup size must match the generated AV shader constant");
+// Below this occupancy the per-entry QK kernel beats the per-row kernel.
+constexpr uint32_t kQkEntryOccupancyFloor = 4096;
+
+uint32_t checked_u32(uint64_t value, const char* label) {
+  if (value > std::numeric_limits<uint32_t>::max()) {
+    throw std::runtime_error(
+        std::string("WebGPU SDPA: ") + label + " exceeds uint32");
   }
-  // out must be [B, H, S_q, D] (same as q's shape).
-  if (out.dims != q.dims) {
+  return static_cast<uint32_t>(value);
+}
+
+uint64_t checked_mul(uint64_t lhs, uint64_t rhs, const char* label) {
+  if (rhs != 0 && lhs > std::numeric_limits<uint64_t>::max() / rhs) {
     throw std::runtime_error(
-        "WebGPU et_vk.sdpa: out shape must match q [B, H, S_q, D]");
+        std::string("WebGPU SDPA: ") + label + " overflow");
   }
+  return lhs * rhs;
+}
 
-  const bool has_mask =
-      graph.get_value_type(mask_id) == WebGPUGraph::ValueType::Tensor;
-  if (has_mask) {
-    // The QK shader indexes mask as [B, H, S_q, S_kv] row-major; require it.
-    const auto& mask = graph.get_tensor(mask_id);
-    if (mask.nbytes != uint64_t(B) * H * S_q * S_kv * sizeof(float)) {
+uint64_t numel(const std::vector<int64_t>& dims, const char* label) {
+  uint64_t value = 1;
+  for (int64_t dim : dims) {
+    if (dim <= 0) {
       throw std::runtime_error(
-          "WebGPU et_vk.sdpa: attn_mask must be [B, H, S_q, S_kv] fp32");
+          std::string("WebGPU SDPA: non-positive ") + label + " dimension");
     }
+    value = checked_mul(value, static_cast<uint64_t>(dim), label);
   }
+  return value;
+}
 
-  float scale = 1.0f / std::sqrt(static_cast<float>(D));
-  const auto scale_type = graph.get_value_type(scale_id);
-  if (scale_type == WebGPUGraph::ValueType::Double) {
-    scale = static_cast<float>(graph.get_double(scale_id));
-  } else if (scale_type != WebGPUGraph::ValueType::Null) {
-    throw std::runtime_error("WebGPU et_vk.sdpa: scale must be Double or None");
+TensorShape parse_shape(
+    const std::vector<int64_t>& dims,
+    SdpaLayout layout,
+    const char* label) {
+  if (dims.size() < 3) {
+    throw std::runtime_error(
+        std::string("WebGPU SDPA: ") + label + " rank must be at least 3");
+  }
+  const size_t rank = dims.size();
+  const size_t h_dim = layout == SdpaLayout::BHSD ? rank - 3 : rank - 2;
+  const size_t s_dim = layout == SdpaLayout::BHSD ? rank - 2 : rank - 3;
+  uint64_t batch = 1;
+  for (size_t i = 0; i + 3 < rank; ++i) {
+    if (dims[i] <= 0) {
+      throw std::runtime_error("WebGPU SDPA: non-positive batch dimension");
+    }
+    batch = checked_mul(batch, static_cast<uint64_t>(dims[i]), "batch");
   }
+  return {
+      checked_u32(batch, "batch"),
+      checked_u32(static_cast<uint64_t>(dims[h_dim]), "heads"),
+      checked_u32(static_cast<uint64_t>(dims[s_dim]), "sequence"),
+      checked_u32(static_cast<uint64_t>(dims[rank - 1]), "head dimension")};
+}
 
-  const uint64_t num_rows = uint64_t(B) * H * S_q; // attn_weights rows
-  const uint64_t aw_numel = num_rows * S_kv; // [B, H, S_q, S_kv]
-  const uint64_t out_numel = uint64_t(B) * H * S_q * D;
-  const size_t aw_bytes = static_cast<size_t>(aw_numel) * sizeof(float);
-
-  // Up-front dispatch-limit checks (throw BEFORE any buffer alloc → no leak).
-  // QK: per-row (one thread per (b,h,s) row, vec4 loads) is fastest for
-  // standard attention, but starves the GPU on channel attention (S_q =
-  // head_dim, so num_rows is tiny → few workgroups serial over a huge S_kv*D).
-  // Route to the per-entry kernel (one thread per (b,h,s,c) attn entry,
-  // 2D-folded) below an occupancy floor; both write a layout-identical
-  // attn[B,H,S_q,S_kv] so softmax/AV are unchanged, and either branch is
-  // numerically correct, so the floor is a perf knob only (Canary M4 Pro:
-  // per-entry ~15-30x faster at num_rows <= 256, per-row wins at num_rows >=
-  // 8192). AV = one per (b,h,s,d4) vec4; softmax = one workgroup per row.
-  constexpr uint32_t kQkEntryOccupancyFloor = 4096u;
-  const bool qk_per_entry = num_rows < kQkEntryOccupancyFloor;
-  const uint32_t qk_wg_size = utils::clamp_workgroup_size(
-      device,
-      qk_per_entry ? kEtVkSdpaQkEntryWorkgroupSizeX
-                   : kEtVkSdpaQkWorkgroupSizeX);
-  uint32_t qk_wg_count = 0;
-  utils::WgCount qk_entry_grid = {};
-  if (qk_per_entry) {
-    qk_entry_grid = utils::compute_2d_workgroup_count(
-        device,
-        static_cast<uint32_t>(aw_numel),
-        qk_wg_size,
-        "et_vk_sdpa_qk_entry");
-  } else {
-    qk_wg_count = utils::compute_1d_workgroup_count(
-        device, static_cast<uint32_t>(num_rows), qk_wg_size, "et_vk_sdpa_qk");
-  }
-  const uint32_t av_wg_size =
-      utils::clamp_workgroup_size(device, kEtVkSdpaAvWorkgroupSizeX);
-  const uint64_t out_numel4 =
-      out_numel / 4; // exact: D % 4 == 0 (checked above)
-  const uint32_t av_wg_count = utils::compute_1d_workgroup_count(
-      device, static_cast<uint32_t>(out_numel4), av_wg_size, "et_vk_sdpa_av");
-  // Near-square 2D grid of workgroups (1 workgroup = 1 row) past the 65535
-  // per-dimension ceiling; sdpa_softmax.wgsl recovers the flat row index from
-  // @builtin(num_workgroups), so no override constant is needed here.
-  utils::WgCount softmax_grid = utils::compute_2d_workgroup_count(
-      device,
-      static_cast<uint32_t>(num_rows),
-      /*workgroup_size=*/1,
-      "et_vk_sdpa_softmax");
-  utils::check_fp32(out, "et_vk.sdpa", "out");
-
-  // NOTE: graph.create_scratch_buffer allocates a fresh buffer per call (no
-  // pooling), so a multi-layer graph holds 2 × num_layers × aw_bytes live, not
-  // 2 × max(aw_bytes) (e.g. S=1024 × 12 layers ≈ 1.1 GB) — a memory-headroom
-  // gap in the shared allocator, not this op; not a correctness issue.
-  WGPUBuffer attn_buf = graph.create_scratch_buffer(aw_bytes);
-  WGPUBuffer softmax_buf = graph.create_scratch_buffer(aw_bytes);
-
-  // ---- Dispatch 1: QK (per-row for standard attn, per-entry for channel) ----
-  {
-    QkParams p = {};
-    p.B = B;
-    p.H = H;
-    p.S_q = S_q;
-    p.S_kv = S_kv;
-    p.D = D;
-    p.has_mask = has_mask ? 1u : 0u;
-    p.scale = scale;
-    WGPUBuffer uniform_buffer =
-        utils::make_uniform(device, &p, sizeof(QkParams));
-    graph.add_uniform_buffer_bytes(sizeof(QkParams));
-
-    // 4-byte dummy storage to satisfy the mask binding when absent (shader
-    // never reads it under has_mask == 0).
-    utils::OptionalBinding mask = utils::make_optional_binding(
-        device,
-        has_mask,
-        has_mask ? graph.get_tensor(mask_id).buffer : nullptr,
-        has_mask ? graph.get_tensor(mask_id).nbytes : 0);
-
-    WGPUConstantEntry wg_const = utils::make_wg_size_constant(qk_wg_size);
-
-    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
-        device,
-        qk_per_entry ? kEtVkSdpaQkEntryWGSL : kEtVkSdpaQkWGSL,
-        {
-            {0, WGPUBufferBindingType_Storage, attn_buf, aw_bytes},
-            {1, WGPUBufferBindingType_ReadOnlyStorage, q.buffer, q.nbytes},
-            {2, WGPUBufferBindingType_ReadOnlyStorage, k.buffer, k.nbytes},
-            {3,
-             WGPUBufferBindingType_ReadOnlyStorage,
-             mask.buffer,
-             mask.nbytes},
-            {4,
-             WGPUBufferBindingType_Uniform,
-             uniform_buffer,
-             sizeof(QkParams)},
-        },
-        &wg_const,
-        1);
+void check_fp32(const WebGPUTensor& tensor, const char* label) {
+  const uint64_t expected =
+      checked_mul(numel(tensor.dims, label), sizeof(float), label);
+  if (tensor.elem_size != sizeof(float) || tensor.is_int ||
+      tensor.nbytes != expected) {
+    throw std::runtime_error(
+        std::string("WebGPU SDPA: ") + label + " must be fp32");
+  }
+}
 
-    if (qk_per_entry) {
-      graph.add_dispatch_2d(
-          bundle.pipeline, bundle.bind_group, qk_entry_grid.x, qk_entry_grid.y);
-    } else {
-      graph.add_dispatch({bundle.pipeline, bundle.bind_group, qk_wg_count});
-    }
+SdpaShape validate_shapes(
+    const std::vector<int64_t>& q_dims,
+    const std::vector<int64_t>& k_dims,
+    const std::vector<int64_t>& v_dims,
+    SdpaLayout layout) {
+  if (q_dims.size() != k_dims.size() || k_dims.size() != v_dims.size()) {
+    throw std::runtime_error("WebGPU SDPA: q/k/v rank mismatch");
+  }
+  const TensorShape q = parse_shape(q_dims, layout, "q");
+  const TensorShape k = parse_shape(k_dims, layout, "k");
+  const TensorShape v = parse_shape(v_dims, layout, "v");
+  if (q.B != k.B || k.B != v.B) {
+    throw std::runtime_error("WebGPU SDPA: q/k/v batch mismatch");
+  }
+  if (q.D != k.D || k.D != v.D) {
+    throw std::runtime_error("WebGPU SDPA: q/k/v head dimension mismatch");
+  }
+  if (k.H != v.H || k.S != v.S) {
+    throw std::runtime_error("WebGPU SDPA: k/v shape mismatch");
+  }
+  if (q.H % k.H != 0) {
+    throw std::runtime_error("WebGPU SDPA: Hq must be divisible by Hkv");
+  }
+  if (q.D % 4 != 0) {
+    throw std::runtime_error(
+        "WebGPU SDPA: head dimension must be a multiple of 4");
+  }
+  return {q.B, q.H, k.H, q.S, k.S, q.D, q.H / k.H};
+}
 
-    wgpuBufferRelease(uniform_buffer);
-    if (mask.owned_dummy != nullptr) {
-      wgpuBufferRelease(mask.owned_dummy);
+MaskMode validate_mask(
+    WebGPUGraph& graph,
+    int mask_id,
+    const SdpaShape& shape,
+    bool require_rank2,
+    bool live) {
+  using VT = WebGPUGraph::ValueType;
+  const VT type = graph.get_value_type(mask_id);
+  if (type == VT::Null) {
+    if (require_rank2) {
+      throw std::runtime_error(
+          "WebGPU et_vk.gemma4_sdpa requires an additive attn_mask");
+    }
+    return MaskMode::None;
+  }
+  if (type != VT::Tensor) {
+    throw std::runtime_error("WebGPU SDPA: attn_mask must be a tensor or None");
+  }
+  const auto& mask = graph.get_tensor(mask_id);
+  if (!live) {
+    check_fp32(mask, "attn_mask");
+  }
+  const auto& dims = live ? graph.cur_dims(mask_id) : mask.dims;
+  if (dims.size() == 2 && dims[0] == shape.S_q && dims[1] == shape.S_kv) {
+    return MaskMode::Rank2;
+  }
+  if (!require_rank2 && dims.size() == 4 && dims[0] == shape.B &&
+      dims[1] == shape.Hq && dims[2] == shape.S_q && dims[3] == shape.S_kv) {
+    return MaskMode::Expanded;
+  }
+  if (require_rank2) {
+    std::string actual = "[";
+    for (size_t i = 0; i < dims.size(); i++) {
+      actual += (i == 0 ? "" : ",") + std::to_string(dims[i]);
     }
+    actual += "]";
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa: attn_mask " + actual +
+        " must equal [S_q,S_kv]=[" + std::to_string(shape.S_q) + "," +
+        std::to_string(shape.S_kv) + "]");
   }
+  throw std::runtime_error(
+      "WebGPU et_vk.sdpa: attn_mask must be [S_q, S_kv] or "
+      "[B, Hq, S_q, S_kv]");
+}
 
-  // ---- Dispatch 2: softmax over the last dim (reuse sdpa_softmax.wgsl) ----
-  {
-    SoftmaxParams p = {};
-    p.num_rows = static_cast<uint32_t>(num_rows);
-    p.row_width = S_kv;
-    WGPUBuffer uniform_buffer =
-        utils::make_uniform(device, &p, sizeof(SoftmaxParams));
-    graph.add_uniform_buffer_bytes(sizeof(SoftmaxParams));
+LiveState make_live_state(
+    WGPUDevice device,
+    const SdpaShape& shape,
+    MaskMode mask_mode,
+    SdpaLayout layout,
+    float scale,
+    uint32_t qk_wg,
+    uint32_t av_wg,
+    uint32_t bh_lo,
+    uint32_t bh_count,
+    bool allow_masked_qk_elision) {
+  // Grids and the scratch extent cover only this chunk's batch-head pairs.
+  const uint64_t num_rows64 =
+      checked_mul(bh_count, shape.S_q, "attention rows");
+  const uint64_t aw_numel64 =
+      checked_mul(num_rows64, shape.S_kv, "attention elements");
+  const uint64_t out_vec4_64 =
+      checked_mul(num_rows64, shape.D / 4u, "output vec4 elements");
+  const uint64_t qk_tiles64 = checked_mul(
+      bh_count,
+      checked_mul(
+          utils::div_up(shape.S_q, kQkTileM),
+          utils::div_up(shape.S_kv, kQkTileN),
+          "attention qk tiles"),
+      "attention qk tiles");
+  const uint32_t num_rows = checked_u32(num_rows64, "attention rows");
+  const uint32_t aw_numel = checked_u32(aw_numel64, "attention elements");
+  const uint32_t out_vec4 = checked_u32(out_vec4_64, "output vec4 elements");
+  const uint32_t qk_tiles = checked_u32(qk_tiles64, "attention qk tiles");
+
+  LiveState state = {};
+  state.qk = {
+      shape.B,
+      shape.Hq,
+      shape.Hkv,
+      shape.S_q,
+      shape.S_kv,
+      shape.D,
+      shape.g,
+      mask_mode == MaskMode::None ? 0u : 1u,
+      static_cast<uint32_t>(mask_mode),
+      static_cast<uint32_t>(layout),
+      0u,
+      scale};
+  state.av = {
+      shape.B,
+      shape.Hq,
+      shape.Hkv,
+      shape.S_q,
+      shape.S_kv,
+      shape.D,
+      shape.g,
+      static_cast<uint32_t>(layout)};
+  state.qk.bh_lo = bh_lo;
+  state.qk.bh_count = bh_count;
+  state.qk.elide_masked_qk = allow_masked_qk_elision &&
+          mask_mode == MaskMode::Rank2 && layout == SdpaLayout::BSHD &&
+          scale == 1.0f && shape.B == 1u && shape.S_q == 1u && shape.Hq == 8u &&
+          shape.Hkv == 1u && shape.g == 8u &&
+          (shape.D == 256u || shape.D == 512u) && shape.S_kv <= 4096u
+      ? 1u
+      : 0u;
+  state.av.bh_lo = bh_lo;
+  state.av.bh_count = bh_count;
+  state.softmax = {num_rows, shape.S_kv, 0u, 0u};
+  state.qk_row_grid = utils::compute_2d_workgroup_count(
+      device, qk_tiles, qk_wg, "et_vk_sdpa_qk");
+  state.qk_entry_grid = utils::compute_2d_workgroup_count(
+      device, aw_numel, qk_wg, "et_vk_sdpa_qk_entry");
+  state.softmax_grid = utils::compute_2d_workgroup_count(
+      device, num_rows, 1u, "et_vk_sdpa_softmax");
+  state.av_grid = utils::compute_2d_workgroup_count(
+      device, out_vec4, av_wg, "et_vk_sdpa_av");
+  state.use_qk_entry = num_rows < kQkEntryOccupancyFloor;
+  return state;
+}
 
-    // sdpa_softmax.wgsl hardcodes @workgroup_size(64,1,1); no override
-    // constant is needed to decode the near-square 2D grid.
-    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
-        device,
-        kSdpaSoftmaxWGSL,
-        {
-            {0, WGPUBufferBindingType_Storage, softmax_buf, aw_bytes},
-            {1, WGPUBufferBindingType_ReadOnlyStorage, attn_buf, aw_bytes},
-            {2,
-             WGPUBufferBindingType_Uniform,
-             uniform_buffer,
-             sizeof(SoftmaxParams)},
-        });
+WGPUConstantEntry workgroup_constant(uint32_t size) {
+  WGPUConstantEntry entry = {};
+  entry.key = {"wg_size", WGPU_STRLEN};
+  entry.value = static_cast<double>(size);
+  return entry;
+}
 
-    graph.add_dispatch_2d(
-        bundle.pipeline, bundle.bind_group, softmax_grid.x, softmax_grid.y);
+size_t record_qk_dispatch(
+    WebGPUGraph& graph,
+    const char* shader,
+    const char* label,
+    const WebGPUTensor& q,
+    const WebGPUTensor& k,
+    WGPUBuffer mask_buffer,
+    uint64_t mask_nbytes,
+    WGPUBuffer attn_buffer,
+    uint64_t aw_bytes,
+    uint64_t attn_offset,
+    WGPUBuffer params_buffer,
+    uint32_t qk_wg,
+    utils::WgCount grid) {
+  const WGPUConstantEntry constant = workgroup_constant(qk_wg);
+  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+      graph.device(),
+      shader,
+      {
+          {0,
+           WGPUBufferBindingType_Storage,
+           attn_buffer,
+           aw_bytes,
+           attn_offset},
+          {1, WGPUBufferBindingType_ReadOnlyStorage, q.buffer, q.nbytes},
+          {2, WGPUBufferBindingType_ReadOnlyStorage, k.buffer, k.nbytes},
+          {3, WGPUBufferBindingType_ReadOnlyStorage, mask_buffer, mask_nbytes},
+          {4, WGPUBufferBindingType_Uniform, params_buffer, sizeof(QkParams)},
+      },
+      &constant,
+      1);
+  return graph.add_dispatch(
+      {bundle.pipeline, bundle.bind_group, grid.x, label, grid.y});
+}
 
-    wgpuBufferRelease(uniform_buffer);
+void rewrite_live_state(
+    WebGPUGraph& graph,
+    const LiveState& state,
+    WGPUBuffer qk_params,
+    WGPUBuffer softmax_params,
+    WGPUBuffer av_params,
+    size_t qk_dispatch,
+    bool fixed_qk_entry,
+    bool dual_qk,
+    size_t qk_route_group,
+    size_t softmax_dispatch,
+    size_t av_dispatch) {
+  wgpuQueueWriteBuffer(
+      graph.queue(), qk_params, 0, &state.qk, sizeof(state.qk));
+  wgpuQueueWriteBuffer(
+      graph.queue(), softmax_params, 0, &state.softmax, sizeof(state.softmax));
+  wgpuQueueWriteBuffer(
+      graph.queue(), av_params, 0, &state.av, sizeof(state.av));
+
+  if (dual_qk) {
+    graph.select_dispatch_route(
+        qk_route_group,
+        state.use_qk_entry ? 1u : 0u,
+        {state.use_qk_entry ? state.qk_entry_grid : state.qk_row_grid});
+  } else {
+    const utils::WgCount grid =
+        fixed_qk_entry ? state.qk_entry_grid : state.qk_row_grid;
+    graph.dispatch_at(qk_dispatch).workgroup_count_x = grid.x;
+    graph.dispatch_at(qk_dispatch).workgroup_count_y = grid.y;
   }
+  graph.dispatch_at(softmax_dispatch).workgroup_count_x = state.softmax_grid.x;
+  graph.dispatch_at(softmax_dispatch).workgroup_count_y = state.softmax_grid.y;
+  graph.dispatch_at(av_dispatch).workgroup_count_x = state.av_grid.x;
+  graph.dispatch_at(av_dispatch).workgroup_count_y = state.av_grid.y;
+}
 
-  // ---- Dispatch 3: AV (one thread per (b,h,s,d4) vec4 output element) ----
-  {
-    AvParams p = {};
-    p.B = B;
-    p.H = H;
-    p.S_q = S_q;
-    p.S_kv = S_kv;
-    p.D = D;
-    WGPUBuffer uniform_buffer =
-        utils::make_uniform(device, &p, sizeof(AvParams));
-    graph.add_uniform_buffer_bytes(sizeof(AvParams));
+void build_sdpa(
+    WebGPUGraph& graph,
+    int q_id,
+    int k_id,
+    int v_id,
+    int mask_id,
+    int out_id,
+    SdpaLayout layout,
+    bool require_rank2_mask,
+    float scale,
+    bool allow_masked_qk_elision) {
+  const auto& q = graph.get_tensor(q_id);
+  const auto& k = graph.get_tensor(k_id);
+  const auto& v = graph.get_tensor(v_id);
+  const auto& out = graph.get_tensor(out_id);
+  check_fp32(q, "q");
+  check_fp32(k, "k");
+  check_fp32(v, "v");
+  check_fp32(out, "out");
+  if (out.dims != q.dims) {
+    throw std::runtime_error("WebGPU SDPA: output shape must match q");
+  }
 
-    WGPUConstantEntry wg_const = utils::make_wg_size_constant(av_wg_size);
+  const SdpaShape max_shape = validate_shapes(q.dims, k.dims, v.dims, layout);
+  const MaskMode mask_mode =
+      validate_mask(graph, mask_id, max_shape, require_rank2_mask, false);
+  const uint64_t aw_numel = checked_mul(
+      checked_mul(
+          checked_mul(max_shape.B, max_shape.Hq, "attention elements"),
+          max_shape.S_q,
+          "attention elements"),
+      max_shape.S_kv,
+      "attention elements");
+  checked_u32(aw_numel, "attention elements");
+  WGPULimits limits = {};
+  if (wgpuDeviceGetLimits(graph.device(), &limits) != WGPUStatus_Success ||
+      limits.maxStorageBufferBindingSize == 0 || limits.maxBufferSize == 0) {
+    throw std::runtime_error("WebGPU SDPA: device limits unavailable");
+  }
+  // Report the shape that produced the size: the byte count alone is ambiguous
+  // (it factors several ways) and guessing which tensor it is has already
+  // produced wrong fixes on this backend.
+  auto scratch_shape = [&]() {
+    return " (B=" + std::to_string(max_shape.B) +
+        " Hq=" + std::to_string(max_shape.Hq) +
+        " S_q=" + std::to_string(max_shape.S_q) +
+        " S_kv=" + std::to_string(max_shape.S_kv) + " fp32, x2 buffers)";
+  };
+  // maxStorageBufferBindingSize caps a binding VIEW, not the buffer. The
+  // scratch is [B][Hq][S_q][S_kv], so a range of batch-head pairs is a
+  // contiguous slice; bind one such chunk per dispatch set.
+  const uint64_t bh_bytes =
+      static_cast<uint64_t>(max_shape.S_q) * max_shape.S_kv * sizeof(float);
+  const utils::RowChunking chunking = utils::compute_row_chunking(
+      limits.maxStorageBufferBindingSize,
+      bh_bytes,
+      static_cast<uint64_t>(max_shape.B) * max_shape.Hq,
+      "et_vk_sdpa attention scratch");
+
+  const uint32_t qk_wg =
+      utils::clamp_workgroup_size(graph.device(), kQkWorkgroupSize);
+  const uint32_t av_wg =
+      utils::clamp_workgroup_size(graph.device(), kAvWorkgroupSize);
+
+  // Chunks run back to back in one compute pass (QK -> softmax -> AV per
+  // chunk), and WebGPU orders dispatches within a pass, so one chunk's worth of
+  // scratch is enough for all of them -- allocate that, not the whole B*Hq
+  // extent.
+  const uint64_t chunk_bytes = chunking.rows_per_chunk * bh_bytes;
+  if (chunk_bytes > limits.maxBufferSize ||
+      chunk_bytes > std::numeric_limits<size_t>::max()) {
+    throw std::runtime_error(
+        "WebGPU SDPA: chunked attention scratch is " +
+        std::to_string(chunk_bytes) +
+        " bytes, over the per-buffer allocation limit of " +
+        std::to_string(
+            std::min<uint64_t>(
+                limits.maxBufferSize, std::numeric_limits<size_t>::max())) +
+        scratch_shape());
+  }
+  WGPUBuffer attn_buffer =
+      graph.acquire_scratch(static_cast<size_t>(chunk_bytes));
+  WebGPUGraph::ScopedScratch attn_guard(&graph, attn_buffer);
+  WGPUBuffer softmax_buffer =
+      graph.acquire_scratch(static_cast<size_t>(chunk_bytes));
+  WebGPUGraph::ScopedScratch softmax_guard(&graph, softmax_buffer);
+
+  const bool has_mask = mask_mode != MaskMode::None;
+  WGPUBuffer mask_buffer = has_mask ? graph.get_tensor(mask_id).buffer
+                                    : graph.create_scratch_buffer(16);
+  const uint64_t mask_nbytes =
+      has_mask ? graph.get_tensor(mask_id).nbytes : 16u;
+
+  const bool directly_dynamic_qk = graph.tensor_has_dynamic_dims(q_id) ||
+      graph.tensor_has_dynamic_dims(k_id) ||
+      graph.tensor_has_dynamic_dims(v_id) ||
+      (has_mask && graph.tensor_has_dynamic_dims(mask_id));
+  const bool unsafe_masked_qk_alias = has_mask &&
+      (q.buffer == k.buffer || q.buffer == v.buffer || q.buffer == out.buffer ||
+       q.buffer == graph.get_tensor(mask_id).buffer || k.buffer == v.buffer ||
+       k.buffer == out.buffer || k.buffer == graph.get_tensor(mask_id).buffer ||
+       v.buffer == out.buffer || v.buffer == graph.get_tensor(mask_id).buffer ||
+       out.buffer == graph.get_tensor(mask_id).buffer);
+  const bool masked_qk_elision_enabled =
+      allow_masked_qk_elision && has_mask && !unsafe_masked_qk_alias;
+
+  const uint64_t bh_total = static_cast<uint64_t>(max_shape.B) * max_shape.Hq;
+  std::vector<ChunkRecord> chunks;
+  for (uint32_t c = 0; c < chunking.num_chunks; c++) {
+    const uint64_t bh_lo = static_cast<uint64_t>(c) * chunking.rows_per_chunk;
+    const uint64_t bh_n = std::min(chunking.rows_per_chunk, bh_total - bh_lo);
+    const uint64_t off = 0; // shared one-chunk scratch, reused per chunk
+    const uint64_t span = bh_n * bh_bytes;
+    const LiveState st = make_live_state(
+        graph.device(),
+        max_shape,
+        mask_mode,
+        layout,
+        scale,
+        qk_wg,
+        av_wg,
+        static_cast<uint32_t>(bh_lo),
+        static_cast<uint32_t>(bh_n),
+        masked_qk_elision_enabled);
+
+    ChunkRecord r = {};
+    r.bh_lo = static_cast<uint32_t>(bh_lo);
+    r.bh_count = static_cast<uint32_t>(bh_n);
+    r.qk_params = graph.make_uniform_buffer(&st.qk, sizeof(st.qk));
+    r.softmax_params =
+        graph.make_uniform_buffer(&st.softmax, sizeof(st.softmax));
+    r.av_params = graph.make_uniform_buffer(&st.av, sizeof(st.av));
+    graph.own_uniform_buffer(r.qk_params);
+    graph.own_uniform_buffer(r.softmax_params);
+    graph.own_uniform_buffer(r.av_params);
+
+    const bool exact_live_entry_route =
+        graph.has_dynamic_shapes() && !directly_dynamic_qk && !st.use_qk_entry;
+    r.dual_qk = directly_dynamic_qk || exact_live_entry_route;
+    r.fixed_qk_entry = st.use_qk_entry;
+    r.qk_dispatch = 0;
+    r.qk_route_group = 0;
+    if (r.dual_qk) {
+      const size_t row_dispatch = record_qk_dispatch(
+          graph,
+          kEtVkSdpaQkWGSL,
+          "et_vk_sdpa_qk",
+          q,
+          k,
+          mask_buffer,
+          mask_nbytes,
+          attn_buffer,
+          span,
+          off,
+          r.qk_params,
+          qk_wg,
+          st.qk_row_grid);
+      const size_t entry_dispatch = record_qk_dispatch(
+          graph,
+          exact_live_entry_route ? kEtVkSdpaQkEntryExactWGSL
+                                 : kEtVkSdpaQkEntryWGSL,
+          exact_live_entry_route ? "et_vk_sdpa_qk_entry_exact"
+                                 : "et_vk_sdpa_qk_entry",
+          q,
+          k,
+          mask_buffer,
+          mask_nbytes,
+          attn_buffer,
+          span,
+          off,
+          r.qk_params,
+          qk_wg,
+          st.qk_entry_grid);
+      r.qk_route_group = graph.register_dispatch_route_group(
+          {{row_dispatch, row_dispatch + 1},
+           {entry_dispatch, entry_dispatch + 1}});
+      graph.select_dispatch_route(
+          r.qk_route_group,
+          st.use_qk_entry ? 1u : 0u,
+          {st.use_qk_entry ? st.qk_entry_grid : st.qk_row_grid});
+    } else {
+      r.qk_dispatch = record_qk_dispatch(
+          graph,
+          st.use_qk_entry ? kEtVkSdpaQkEntryWGSL : kEtVkSdpaQkWGSL,
+          st.use_qk_entry ? "et_vk_sdpa_qk_entry" : "et_vk_sdpa_qk",
+          q,
+          k,
+          mask_buffer,
+          mask_nbytes,
+          attn_buffer,
+          span,
+          off,
+          r.qk_params,
+          qk_wg,
+          st.use_qk_entry ? st.qk_entry_grid : st.qk_row_grid);
+    }
 
-    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
-        device,
+    utils::ComputePipelineBundle sm_pipeline = utils::make_compute_pipeline(
+        graph.device(),
+        kSdpaSoftmaxWGSL,
+        {
+            {0, WGPUBufferBindingType_Storage, softmax_buffer, span, off},
+            {1, WGPUBufferBindingType_ReadOnlyStorage, attn_buffer, span, off},
+            {2,
+             WGPUBufferBindingType_Uniform,
+             r.softmax_params,
+             sizeof(SoftmaxParams)},
+        });
+    r.softmax_dispatch = graph.add_dispatch(
+        {sm_pipeline.pipeline,
+         sm_pipeline.bind_group,
+         st.softmax_grid.x,
+         "et_vk_sdpa_softmax",
+         st.softmax_grid.y});
+
+    const WGPUConstantEntry av_constant = workgroup_constant(av_wg);
+    utils::ComputePipelineBundle av_pipeline = utils::make_compute_pipeline(
+        graph.device(),
         kEtVkSdpaAvWGSL,
         {
             {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
-            {1, WGPUBufferBindingType_ReadOnlyStorage, softmax_buf, aw_bytes},
+            {1,
+             WGPUBufferBindingType_ReadOnlyStorage,
+             softmax_buffer,
+             span,
+             off},
             {2, WGPUBufferBindingType_ReadOnlyStorage, v.buffer, v.nbytes},
-            {3,
-             WGPUBufferBindingType_Uniform,
-             uniform_buffer,
-             sizeof(AvParams)},
+            {3, WGPUBufferBindingType_Uniform, r.av_params, sizeof(AvParams)},
         },
-        &wg_const,
+        &av_constant,
         1);
+    r.av_dispatch = graph.add_dispatch(
+        {av_pipeline.pipeline,
+         av_pipeline.bind_group,
+         st.av_grid.x,
+         "et_vk_sdpa_av",
+         st.av_grid.y});
+    chunks.push_back(r);
+  }
 
-    graph.add_dispatch({bundle.pipeline, bundle.bind_group, av_wg_count});
+  auto resize = [q_id,
+                 k_id,
+                 v_id,
+                 mask_id,
+                 out_id,
+                 layout,
+                 require_rank2_mask,
+                 max_shape,
+                 scale,
+                 qk_wg,
+                 av_wg,
+                 masked_qk_elision_enabled,
+                 chunks](WebGPUGraph& gr) {
+    const SdpaShape live_shape = validate_shapes(
+        gr.cur_dims(q_id), gr.cur_dims(k_id), gr.cur_dims(v_id), layout);
+    if (live_shape.B != max_shape.B || live_shape.Hq != max_shape.Hq ||
+        live_shape.Hkv != max_shape.Hkv || live_shape.D != max_shape.D ||
+        live_shape.S_q > max_shape.S_q || live_shape.S_kv > max_shape.S_kv) {
+      throw std::runtime_error(
+          "WebGPU SDPA: live shape exceeds allocation bounds");
+    }
+    const MaskMode live_mask =
+        validate_mask(gr, mask_id, live_shape, require_rank2_mask, true);
+    // The batch-head split is shape-independent (B and Hq cannot change here),
+    // so each chunk keeps its range and only its grids/params are rewritten.
+    for (const ChunkRecord& r : chunks) {
+      const LiveState state = make_live_state(
+          gr.device(),
+          live_shape,
+          live_mask,
+          layout,
+          scale,
+          qk_wg,
+          av_wg,
+          r.bh_lo,
+          r.bh_count,
+          masked_qk_elision_enabled);
+      rewrite_live_state(
+          gr,
+          state,
+          r.qk_params,
+          r.softmax_params,
+          r.av_params,
+          r.qk_dispatch,
+          r.fixed_qk_entry,
+          r.dual_qk,
+          r.qk_route_group,
+          r.softmax_dispatch,
+          r.av_dispatch);
+    }
+    gr.set_cur_dims(out_id, gr.cur_dims(q_id));
+  };
+  graph.add_tensor_resize_hook(q_id, resize);
+  graph.add_tensor_resize_hook(k_id, resize);
+  graph.add_tensor_resize_hook(v_id, resize);
+  if (has_mask) {
+    graph.add_tensor_resize_hook(mask_id, resize);
+  }
+}
 
-    wgpuBufferRelease(uniform_buffer);
+void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  if (args.size() != 6) {
+    throw std::runtime_error("WebGPU et_vk.sdpa: expected 6 args");
+  }
+  const int q_id = args.at(0);
+  const int k_id = args.at(1);
+  const int v_id = args.at(2);
+  const int mask_id = args.at(3);
+  const int scale_id = args.at(4);
+  const int out_id = args.at(5);
+
+  const auto& q = graph.get_tensor(q_id);
+  const TensorShape q_shape = parse_shape(q.dims, SdpaLayout::BHSD, "q");
+  float scale = 1.0f / std::sqrt(static_cast<float>(q_shape.D));
+  const auto scale_type = graph.get_value_type(scale_id);
+  if (scale_type == WebGPUGraph::ValueType::Double) {
+    scale = static_cast<float>(graph.get_double(scale_id));
+  } else if (scale_type != WebGPUGraph::ValueType::Null) {
+    throw std::runtime_error("WebGPU et_vk.sdpa: scale must be Double or None");
+  }
+  if (!std::isfinite(scale)) {
+    throw std::runtime_error("WebGPU et_vk.sdpa: scale must be finite");
+  }
+  build_sdpa(
+      graph,
+      q_id,
+      k_id,
+      v_id,
+      mask_id,
+      out_id,
+      SdpaLayout::BHSD,
+      false,
+      scale,
+      false);
+}
+
+void gemma4_sdpa_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  // Fence: only the Gemma4 exporter's masked, non-causal, scale=1.0 ABI.
+  if (args.size() != 9) {
+    throw std::runtime_error("WebGPU et_vk.gemma4_sdpa: expected 9 args");
+  }
+  const int q_id = args.at(0);
+  const int k_id = args.at(1);
+  const int v_id = args.at(2);
+  const int start_id = args.at(3);
+  const int mask_id = args.at(4);
+  const int dropout_id = args.at(5);
+  const int causal_id = args.at(6);
+  const int scale_id = args.at(7);
+  const int out_id = args.at(8);
+  using VT = WebGPUGraph::ValueType;
+
+  const VT start_type = graph.get_value_type(start_id);
+  int64_t start = 0;
+  if (start_type == VT::Int) {
+    start = graph.get_int(start_id);
+  } else if (start_type == VT::SymInt) {
+    start = graph.read_symint(start_id);
+  } else {
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa: start_pos must be Int or SymInt");
+  }
+  if (start < 0) {
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa: start_pos must be non-negative");
+  }
+
+  const VT dropout_type = graph.get_value_type(dropout_id);
+  const bool dropout_zero = dropout_type == VT::Null ||
+      (dropout_type == VT::Double && graph.get_double(dropout_id) == 0.0) ||
+      (dropout_type == VT::Int && graph.get_int(dropout_id) == 0);
+  if (!dropout_zero) {
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa: only dropout_p=0 is supported");
+  }
+  const VT causal_type = graph.get_value_type(causal_id);
+  const bool causal_false = causal_type == VT::Null ||
+      (causal_type == VT::Bool && !graph.get_bool(causal_id));
+  if (!causal_false) {
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa: only is_causal=false is supported");
+  }
+  if (graph.get_value_type(mask_id) != VT::Tensor) {
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa requires an additive attn_mask");
+  }
+  const VT scale_type = graph.get_value_type(scale_id);
+  // Vulkan interning can serialize semantic 1.0 as Int(1); accept both.
+  double scale_value;
+  if (scale_type == VT::Double) {
+    scale_value = graph.get_double(scale_id);
+  } else if (scale_type == VT::Int) {
+    scale_value = static_cast<double>(graph.get_int(scale_id));
+  } else {
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa: requires explicit scale=1.0");
+  }
+  if (!std::isfinite(scale_value)) {
+    throw std::runtime_error("WebGPU et_vk.gemma4_sdpa: scale must be finite");
+  }
+  if (scale_value != 1.0) {
+    throw std::runtime_error(
+        "WebGPU et_vk.gemma4_sdpa: requires explicit scale=1.0");
+  }
+  const float scale = static_cast<float>(scale_value);
+
+  build_sdpa(
+      graph,
+      q_id,
+      k_id,
+      v_id,
+      mask_id,
+      out_id,
+      SdpaLayout::BSHD,
+      true,
+      scale,
+      true);
+  if (start_type == VT::SymInt) {
+    graph.add_resize_hook(start_id, [start_id](WebGPUGraph& gr) {
+      if (gr.read_symint(start_id) < 0) {
+        throw std::runtime_error(
+            "WebGPU et_vk.gemma4_sdpa: start_pos must be non-negative");
+      }
+    });
   }
 }
 
@@ -324,6 +871,7 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
 WEBGPU_REGISTER_OPERATORS {
   WEBGPU_REGISTER_OP(et_vk.sdpa.default, et_vk_sdpa_impl);
+  WEBGPU_REGISTER_OP(et_vk.gemma4_sdpa.default, gemma4_sdpa_impl);
 }
 
 } // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl
index 3ccd1d200a2..86ee5657428 100644
--- a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl
@@ -4,43 +4,59 @@
 
 struct Params {
   B: u32,
-  H: u32,
+  Hq: u32,
+  Hkv: u32,
   S_q: u32,
   S_kv: u32,
   D: u32,
+  g: u32,
+  tensor_layout: u32,
+  bh_lo: u32,
+  bh_count: u32,
   _pad0: u32,
   _pad1: u32,
-  _pad2: u32,
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
 override wg_size: u32 = 64;
 
-// Non-causal fused SDPA, AV phase. out[b,h,s,d]=sum_c sm[b,h,s,c]*v[b,h,c,d].
-// DSHB layout, row-major: out/sm rows over S_q, v rows over S_kv; v/out viewed
-// as vec4<f32> over D (caller guarantees D % 4 == 0 for every model in scope).
-// Supports asymmetric seq (S_q != S_kv); reduces to self-attention when
-// S_q == S_kv. ONE thread per (b, h, s, d4) computing 4 output elements; the
-// thread contracts over c (0..S_kv, scalar — S_kv isn't guaranteed % 4 == 0).
+fn v_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * d4_count;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * d4_count;
+}
+
+fn out_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * d4_count;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * d4_count;
+}
+
 @compute @workgroup_size(wg_size)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
   let d4_count = params.D / 4u;
-  let total = params.B * params.H * params.S_q * d4_count;
-  let i = gid.x;
-  if (i >= total) {
+  let chunk_total = params.bh_count * params.S_q * d4_count;
+  let tid = gid.x + gid.y * (nwg.x * wg_size);
+  if (tid >= chunk_total) {
     return;
   }
+  let i = params.bh_lo * params.S_q * d4_count + tid;
   let d4 = i % d4_count;
   let s = (i / d4_count) % params.S_q;
-  let h = (i / (d4_count * params.S_q)) % params.H;
-  let b = i / (d4_count * params.S_q * params.H);
-
-  let smbase = ((b * params.H + h) * params.S_q + s) * params.S_kv;
-  let vblock4 = (b * params.H + h) * params.S_kv; // first V row of this (b, h)
+  let h = (i / (d4_count * params.S_q)) % params.Hq;
+  let b = i / (d4_count * params.S_q * params.Hq);
+  let kv_h = h / params.g;
+  let smbase =
+      ((b * params.Hq + h - params.bh_lo) * params.S_q + s) * params.S_kv;
 
   var acc: vec4<f32> = vec4<f32>(0.0);
   for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) {
-    acc = acc + sm[smbase + c] * v[(vblock4 + c) * d4_count + d4];
+    acc = acc +
+        sm[smbase + c] * v[v_row4(b, kv_h, c, d4_count) + d4];
   }
-  out[i] = acc;
+  out[out_row4(b, h, s, d4_count) + d4] = acc;
 }
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h
index 6f625437589..a843924a060 100644
--- a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from et_vk_sdpa_av.wgsl - DO NOT EDIT.
-// wgsl-sha256: 0f91462226b666d2fc05381866caed709a6ef3ed2f2562addfe9a23ac81f173b
+// wgsl-sha256: 89cc190126bbf816145901db7a09451d0ef571a0bdccb023ee478a6c7305290d
 inline constexpr const char* kEtVkSdpaAvWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> out: array<vec4<f32>>;
 @group(0) @binding(1) var<storage, read> sm: array<f32>;
@@ -21,45 +21,61 @@ inline constexpr const char* kEtVkSdpaAvWGSL = R"(
 
 struct Params {
   B: u32,
-  H: u32,
+  Hq: u32,
+  Hkv: u32,
   S_q: u32,
   S_kv: u32,
   D: u32,
+  g: u32,
+  tensor_layout: u32,
+  bh_lo: u32,
+  bh_count: u32,
   _pad0: u32,
   _pad1: u32,
-  _pad2: u32,
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
 override wg_size: u32 = 64;
 
-// Non-causal fused SDPA, AV phase. out[b,h,s,d]=sum_c sm[b,h,s,c]*v[b,h,c,d].
-// DSHB layout, row-major: out/sm rows over S_q, v rows over S_kv; v/out viewed
-// as vec4<f32> over D (caller guarantees D % 4 == 0 for every model in scope).
-// Supports asymmetric seq (S_q != S_kv); reduces to self-attention when
-// S_q == S_kv. ONE thread per (b, h, s, d4) computing 4 output elements; the
-// thread contracts over c (0..S_kv, scalar — S_kv isn't guaranteed % 4 == 0).
+fn v_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * d4_count;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * d4_count;
+}
+
+fn out_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * d4_count;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * d4_count;
+}
+
 @compute @workgroup_size(wg_size)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
   let d4_count = params.D / 4u;
-  let total = params.B * params.H * params.S_q * d4_count;
-  let i = gid.x;
-  if (i >= total) {
+  let chunk_total = params.bh_count * params.S_q * d4_count;
+  let tid = gid.x + gid.y * (nwg.x * wg_size);
+  if (tid >= chunk_total) {
     return;
   }
+  let i = params.bh_lo * params.S_q * d4_count + tid;
   let d4 = i % d4_count;
   let s = (i / d4_count) % params.S_q;
-  let h = (i / (d4_count * params.S_q)) % params.H;
-  let b = i / (d4_count * params.S_q * params.H);
-
-  let smbase = ((b * params.H + h) * params.S_q + s) * params.S_kv;
-  let vblock4 = (b * params.H + h) * params.S_kv; // first V row of this (b, h)
+  let h = (i / (d4_count * params.S_q)) % params.Hq;
+  let b = i / (d4_count * params.S_q * params.Hq);
+  let kv_h = h / params.g;
+  let smbase =
+      ((b * params.Hq + h - params.bh_lo) * params.S_q + s) * params.S_kv;
 
   var acc: vec4<f32> = vec4<f32>(0.0);
   for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) {
-    acc = acc + sm[smbase + c] * v[(vblock4 + c) * d4_count + d4];
+    acc = acc +
+        sm[smbase + c] * v[v_row4(b, kv_h, c, d4_count) + d4];
   }
-  out[i] = acc;
+  out[out_row4(b, h, s, d4_count) + d4] = acc;
 }
 )";
 
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl
index f1c4f29ed8c..71e2a2243e7 100644
--- a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl
@@ -5,51 +5,136 @@
 
 struct Params {
   B: u32,
-  H: u32,
+  Hq: u32,
+  Hkv: u32,
   S_q: u32,
   S_kv: u32,
   D: u32,
+  g: u32,
   has_mask: u32,
+  mask_mode: u32,
+  tensor_layout: u32,
   _pad0: u32,
   scale: f32,
+  bh_lo: u32,
+  bh_count: u32,
+  _pad1: u32,
+  _pad2: u32,
 }
 @group(0) @binding(4) var<uniform> params: Params;
 
 override wg_size: u32 = 64;
 
-// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D],
-// k [B, H, S_kv, D]; q/k viewed as vec4<f32> over D (caller guarantees
-// D % 4 == 0 for every model in scope). Supports asymmetric seq (S_q != S_kv,
-// e.g. Hiera pooled query); when S_q == S_kv this reduces to plain
-// self-attention. ONE thread per (b, h, s) ROW of attn_weights
-// [B, H, S_q, S_kv]; the thread loops over c (0..S_kv) and d4 (0..D/4). Row
-// count = B*H*S_q stays well under the 65535 1D dispatch limit for any ViT.
+// TMxTN register tile; mirrors sdpa/sdpa_compute_attn_weights.wgsl (not a knob).
+const TM: u32 = 8u;
+const TN: u32 = 4u;
+
+fn q_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * d4_count;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * d4_count;
+}
+
+fn k_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * d4_count;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * d4_count;
+}
+
+fn load_q4(b: u32, h: u32, s: u32, d4: u32, d4_count: u32) -> vec4<f32> {
+  if (s >= params.S_q) {
+    return vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+  return q[q_row4(b, h, s, d4_count) + d4];
+}
+
+fn load_k4(b: u32, h: u32, c: u32, d4: u32, d4_count: u32) -> vec4<f32> {
+  if (c >= params.S_kv) {
+    return vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+  return k[k_row4(b, h, c, d4_count) + d4];
+}
+
+fn store_qk(
+    scratch_row: u32,
+    absolute_row: u32,
+    s: u32,
+    c: u32,
+    raw: f32) {
+  if (c >= params.S_kv) {
+    return;
+  }
+  var val = raw * params.scale;
+  if (params.has_mask != 0u) {
+    if (params.mask_mode == 1u) {
+      val = val + mask[params.S_kv * s + c];
+    } else {
+      val = val + mask[absolute_row + c];
+    }
+  }
+  attn[scratch_row + c] = val;
+}
+
 @compute @workgroup_size(wg_size)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let num_rows = params.B * params.H * params.S_q;
-  let row = gid.x;
-  if (row >= num_rows) {
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  let nrt = (params.S_q + TM - 1u) / TM;
+  let nct = (params.S_kv + TN - 1u) / TN;
+  let tiles = nrt * nct;
+  let total = tiles * params.bh_count;
+  // 2D dispatch fold: recover the linear tile index across x/y.
+  let idx = gid.x + gid.y * (nwg.x * wg_size);
+  if (idx >= total) {
     return;
   }
-  let s = row % params.S_q;
-  let h = (row / params.S_q) % params.H;
-  let b = row / (params.S_q * params.H);
 
+  // Tile within one (b, h) so a tile never straddles a head boundary.
+  let local_bh = idx / tiles;
+  let bh = params.bh_lo + local_bh;
+  let rem = idx % tiles;
+  let h = bh % params.Hq;
+  let b = bh / params.Hq;
+  let s0 = (rem / nct) * TM;
+  let c0 = (rem % nct) * TN;
+  let kv_h = h / params.g;
   let d4_count = params.D / 4u;
-  let qbase4 = ((b * params.H + h) * params.S_q + s) * d4_count;
-  let kblock4 = (b * params.H + h) * params.S_kv; // first K row of this (b, h)
-  let arow = ((b * params.H + h) * params.S_q + s) * params.S_kv;
-
-  for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) {
-    let kbase4 = (kblock4 + c) * d4_count;
-    var acc: f32 = 0.0;
-    for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
-      acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]);
+
+  var acc: array<vec4<f32>, TM>;
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    acc[i] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+
+  for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
+    var qv: array<vec4<f32>, TM>;
+    var kv: array<vec4<f32>, TN>;
+    for (var i: u32 = 0u; i < TM; i = i + 1u) {
+      qv[i] = load_q4(b, h, s0 + i, d4, d4_count);
+    }
+    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+      kv[j] = load_k4(b, kv_h, c0 + j, d4, d4_count);
     }
-    acc = acc * params.scale;
-    if (params.has_mask != 0u) {
-      acc = acc + mask[arow + c];
+    for (var i: u32 = 0u; i < TM; i = i + 1u) {
+      acc[i] = acc[i] +
+          vec4<f32>(
+              dot(qv[i], kv[0]),
+              dot(qv[i], kv[1]),
+              dot(qv[i], kv[2]),
+              dot(qv[i], kv[3]));
+    }
+  }
+
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    let s = s0 + i;
+    if (s < params.S_q) {
+      let absolute_row = (bh * params.S_q + s) * params.S_kv;
+      let scratch_row = (local_bh * params.S_q + s) * params.S_kv;
+      let av = acc[i];
+      store_qk(scratch_row, absolute_row, s, c0 + 0u, av.x);
+      store_qk(scratch_row, absolute_row, s, c0 + 1u, av.y);
+      store_qk(scratch_row, absolute_row, s, c0 + 2u, av.z);
+      store_qk(scratch_row, absolute_row, s, c0 + 3u, av.w);
     }
-    attn[arow + c] = acc;
   }
 }
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl
index 605ac52c89e..528c2cd2eff 100644
--- a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl
@@ -5,50 +5,70 @@
 
 struct Params {
   B: u32,
-  H: u32,
+  Hq: u32,
+  Hkv: u32,
   S_q: u32,
   S_kv: u32,
   D: u32,
+  g: u32,
   has_mask: u32,
+  mask_mode: u32,
+  tensor_layout: u32,
   _pad0: u32,
   scale: f32,
+  bh_lo: u32,
+  bh_count: u32,
+  _pad1: u32,
+  _pad2: u32,
 }
 @group(0) @binding(4) var<uniform> params: Params;
 
 override wg_size: u32 = 64;
 
-// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D],
-// k [B, H, S_kv, D]. ONE thread per ENTRY (b,h,s,c) of attn_weights
-// [B, H, S_q, S_kv] = one D-length dot. Previously this was one thread per ROW
-// (looping c) — fine for window/self attention but catastrophic for DaViT
-// CHANNEL attention where S_q = head_dim (~32), so B*H*S_q was 128/256/512/1024
-// → only 2/4/8/16 workgroups serial over the huge spatial D (the (2,1,1)@103ms
-// dispatch). Parallelizing over all B*H*S_q*S_kv entries (2D-folded past the
-// 65535 ceiling, mirroring the softmax phase) gives S_kv× more threads.
+fn q_row(b: u32, h: u32, s: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * params.D;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * params.D;
+}
+
+fn k_row(b: u32, h: u32, s: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * params.D;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * params.D;
+}
+
 @compute @workgroup_size(wg_size)
 fn main(
     @builtin(global_invocation_id) gid: vec3<u32>,
     @builtin(num_workgroups) nwg: vec3<u32>) {
-  let aw_numel = params.B * params.H * params.S_q * params.S_kv;
-  let idx = gid.x + gid.y * (nwg.x * wg_size); // 2D-folded linear entry id
-  if (idx >= aw_numel) {
+  let chunk_numel = params.bh_count * params.S_q * params.S_kv;
+  let tid = gid.x + gid.y * (nwg.x * wg_size);
+  if (tid >= chunk_numel) {
     return;
   }
+  let idx = params.bh_lo * params.S_q * params.S_kv + tid;
   let c = idx % params.S_kv;
-  let row = idx / params.S_kv; // (b,h,s) flattened
+  let row = idx / params.S_kv;
   let s = row % params.S_q;
-  let h = (row / params.S_q) % params.H;
-  let b = row / (params.S_q * params.H);
+  let h = (row / params.S_q) % params.Hq;
+  let b = row / (params.S_q * params.Hq);
+  let kv_h = h / params.g;
+  let qbase = q_row(b, h, s);
+  let kbase = k_row(b, kv_h, c);
 
-  let qbase = ((b * params.H + h) * params.S_q + s) * params.D;
-  let kbase = ((b * params.H + h) * params.S_kv + c) * params.D;
   var acc: f32 = 0.0;
   for (var d: u32 = 0u; d < params.D; d = d + 1u) {
     acc = acc + q[qbase + d] * k[kbase + d];
   }
   acc = acc * params.scale;
   if (params.has_mask != 0u) {
-    acc = acc + mask[idx];
+    if (params.mask_mode == 1u) {
+      acc = acc + mask[params.S_kv * s + c];
+    } else {
+      acc = acc + mask[idx];
+    }
   }
-  attn[idx] = acc; // attn is [B,H,S_q,S_kv] row-major -> index == idx
+  attn[tid] = acc;
 }
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact.wgsl
new file mode 100644
index 00000000000..f8a20b5f5e8
--- /dev/null
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact.wgsl
@@ -0,0 +1,108 @@
+// 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> attn: array<f32>;
+@group(0) @binding(1) var<storage, read> q: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> k: array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read> mask: array<f32>;
+
+struct Params {
+  B: u32,
+  Hq: u32,
+  Hkv: u32,
+  S_q: u32,
+  S_kv: u32,
+  D: u32,
+  g: u32,
+  has_mask: u32,
+  mask_mode: u32,
+  tensor_layout: u32,
+  _pad0: u32,
+  scale: f32,
+  bh_lo: u32,
+  bh_count: u32,
+  elide_masked_qk: u32,
+  _pad2: u32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64;
+
+fn q_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * d4_count;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * d4_count;
+}
+
+fn k_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * d4_count;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * d4_count;
+}
+
+fn all_finite(value: vec4<f32>) -> bool {
+  let exponent = bitcast<vec4<u32>>(value) & vec4<u32>(0x7f800000u);
+  return all(exponent != vec4<u32>(0x7f800000u));
+}
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
+  let chunk_numel = params.bh_count * params.S_q * params.S_kv;
+  let tid = gid.x + gid.y * (nwg.x * wg_size);
+  if (tid >= chunk_numel) {
+    return;
+  }
+  let idx = params.bh_lo * params.S_q * params.S_kv + tid;
+  let c = idx % params.S_kv;
+  let row = idx / params.S_kv;
+  let s = row % params.S_q;
+  let h = (row / params.S_q) % params.Hq;
+  let b = row / (params.S_q * params.Hq);
+  let kv_h = h / params.g;
+  let d4_count = params.D / 4u;
+  let qbase4 = q_row4(b, h, s, d4_count);
+  let kbase4 = k_row4(b, kv_h, c, d4_count);
+
+  var mask_value: f32 = 0.0;
+  if (params.has_mask != 0u) {
+    if (params.mask_mode == 1u) {
+      mask_value = mask[params.S_kv * s + c];
+    } else {
+      mask_value = mask[idx];
+    }
+  }
+  if (params.elide_masked_qk != 0u && params.mask_mode == 1u &&
+      bitcast<u32>(mask_value) == 0xff800000u) {
+    var finite_inputs = true;
+    for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
+      finite_inputs = finite_inputs && all_finite(q[qbase4 + d4]) &&
+          all_finite(k[kbase4 + d4]);
+    }
+    if (finite_inputs) {
+      attn[tid] = mask_value;
+      return;
+    }
+  }
+
+  var acc: f32 = 0.0;
+  for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
+    acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]);
+  }
+  acc = acc * params.scale;
+  if (params.has_mask != 0u) {
+    acc = acc + mask_value;
+  }
+  attn[tid] = acc;
+}
+// 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.
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact_wgsl.h
new file mode 100644
index 00000000000..58d4f18196e
--- /dev/null
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_exact_wgsl.h
@@ -0,0 +1,132 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from et_vk_sdpa_qk_entry_exact.wgsl - DO NOT EDIT.
+// wgsl-sha256: 29d50971292eec9f5f87c9043ea3c07fd3314fe4355146b8b91856949944ebf6
+inline constexpr const char* kEtVkSdpaQkEntryExactWGSL = R"(
+// 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> attn: array<f32>;
+@group(0) @binding(1) var<storage, read> q: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> k: array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read> mask: array<f32>;
+
+struct Params {
+  B: u32,
+  Hq: u32,
+  Hkv: u32,
+  S_q: u32,
+  S_kv: u32,
+  D: u32,
+  g: u32,
+  has_mask: u32,
+  mask_mode: u32,
+  tensor_layout: u32,
+  _pad0: u32,
+  scale: f32,
+  bh_lo: u32,
+  bh_count: u32,
+  elide_masked_qk: u32,
+  _pad2: u32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64;
+
+fn q_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * d4_count;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * d4_count;
+}
+
+fn k_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * d4_count;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * d4_count;
+}
+
+fn all_finite(value: vec4<f32>) -> bool {
+  let exponent = bitcast<vec4<u32>>(value) & vec4<u32>(0x7f800000u);
+  return all(exponent != vec4<u32>(0x7f800000u));
+}
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
+  let chunk_numel = params.bh_count * params.S_q * params.S_kv;
+  let tid = gid.x + gid.y * (nwg.x * wg_size);
+  if (tid >= chunk_numel) {
+    return;
+  }
+  let idx = params.bh_lo * params.S_q * params.S_kv + tid;
+  let c = idx % params.S_kv;
+  let row = idx / params.S_kv;
+  let s = row % params.S_q;
+  let h = (row / params.S_q) % params.Hq;
+  let b = row / (params.S_q * params.Hq);
+  let kv_h = h / params.g;
+  let d4_count = params.D / 4u;
+  let qbase4 = q_row4(b, h, s, d4_count);
+  let kbase4 = k_row4(b, kv_h, c, d4_count);
+
+  var mask_value: f32 = 0.0;
+  if (params.has_mask != 0u) {
+    if (params.mask_mode == 1u) {
+      mask_value = mask[params.S_kv * s + c];
+    } else {
+      mask_value = mask[idx];
+    }
+  }
+  if (params.elide_masked_qk != 0u && params.mask_mode == 1u &&
+      bitcast<u32>(mask_value) == 0xff800000u) {
+    var finite_inputs = true;
+    for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
+      finite_inputs = finite_inputs && all_finite(q[qbase4 + d4]) &&
+          all_finite(k[kbase4 + d4]);
+    }
+    if (finite_inputs) {
+      attn[tid] = mask_value;
+      return;
+    }
+  }
+
+  var acc: f32 = 0.0;
+  for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
+    acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]);
+  }
+  acc = acc * params.scale;
+  if (params.has_mask != 0u) {
+    acc = acc + mask_value;
+  }
+  attn[tid] = acc;
+}
+// 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.
+)";
+
+inline constexpr uint32_t kEtVkSdpaQkEntryExactWorkgroupSizeX = 64;
+inline constexpr uint32_t kEtVkSdpaQkEntryExactWorkgroupSizeY = 1;
+inline constexpr uint32_t kEtVkSdpaQkEntryExactWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h
index 6918c6e2621..979daeba0ae 100644
--- a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from et_vk_sdpa_qk_entry.wgsl - DO NOT EDIT.
-// wgsl-sha256: 95ff9f19757d23c41d82054b779c02fb8d83531eef3e8cda5f3a78dde42a4e58
+// wgsl-sha256: 9941f3ef67238d4a60dc24c9cc60320b50e0f647bd699b06c3f4fcb4cc0d449e
 inline constexpr const char* kEtVkSdpaQkEntryWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> attn: array<f32>;
 @group(0) @binding(1) var<storage, read> q: array<f32>;
@@ -22,52 +22,72 @@ inline constexpr const char* kEtVkSdpaQkEntryWGSL = R"(
 
 struct Params {
   B: u32,
-  H: u32,
+  Hq: u32,
+  Hkv: u32,
   S_q: u32,
   S_kv: u32,
   D: u32,
+  g: u32,
   has_mask: u32,
+  mask_mode: u32,
+  tensor_layout: u32,
   _pad0: u32,
   scale: f32,
+  bh_lo: u32,
+  bh_count: u32,
+  _pad1: u32,
+  _pad2: u32,
 }
 @group(0) @binding(4) var<uniform> params: Params;
 
 override wg_size: u32 = 64;
 
-// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D],
-// k [B, H, S_kv, D]. ONE thread per ENTRY (b,h,s,c) of attn_weights
-// [B, H, S_q, S_kv] = one D-length dot. Previously this was one thread per ROW
-// (looping c) — fine for window/self attention but catastrophic for DaViT
-// CHANNEL attention where S_q = head_dim (~32), so B*H*S_q was 128/256/512/1024
-// → only 2/4/8/16 workgroups serial over the huge spatial D (the (2,1,1)@103ms
-// dispatch). Parallelizing over all B*H*S_q*S_kv entries (2D-folded past the
-// 65535 ceiling, mirroring the softmax phase) gives S_kv× more threads.
+fn q_row(b: u32, h: u32, s: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * params.D;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * params.D;
+}
+
+fn k_row(b: u32, h: u32, s: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * params.D;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * params.D;
+}
+
 @compute @workgroup_size(wg_size)
 fn main(
     @builtin(global_invocation_id) gid: vec3<u32>,
     @builtin(num_workgroups) nwg: vec3<u32>) {
-  let aw_numel = params.B * params.H * params.S_q * params.S_kv;
-  let idx = gid.x + gid.y * (nwg.x * wg_size); // 2D-folded linear entry id
-  if (idx >= aw_numel) {
+  let chunk_numel = params.bh_count * params.S_q * params.S_kv;
+  let tid = gid.x + gid.y * (nwg.x * wg_size);
+  if (tid >= chunk_numel) {
     return;
   }
+  let idx = params.bh_lo * params.S_q * params.S_kv + tid;
   let c = idx % params.S_kv;
-  let row = idx / params.S_kv; // (b,h,s) flattened
+  let row = idx / params.S_kv;
   let s = row % params.S_q;
-  let h = (row / params.S_q) % params.H;
-  let b = row / (params.S_q * params.H);
+  let h = (row / params.S_q) % params.Hq;
+  let b = row / (params.S_q * params.Hq);
+  let kv_h = h / params.g;
+  let qbase = q_row(b, h, s);
+  let kbase = k_row(b, kv_h, c);
 
-  let qbase = ((b * params.H + h) * params.S_q + s) * params.D;
-  let kbase = ((b * params.H + h) * params.S_kv + c) * params.D;
   var acc: f32 = 0.0;
   for (var d: u32 = 0u; d < params.D; d = d + 1u) {
     acc = acc + q[qbase + d] * k[kbase + d];
   }
   acc = acc * params.scale;
   if (params.has_mask != 0u) {
-    acc = acc + mask[idx];
+    if (params.mask_mode == 1u) {
+      acc = acc + mask[params.S_kv * s + c];
+    } else {
+      acc = acc + mask[idx];
+    }
   }
-  attn[idx] = acc; // attn is [B,H,S_q,S_kv] row-major -> index == idx
+  attn[tid] = acc;
 }
 )";
 
diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h
index 030e9dec24e..eecd91a4d8b 100644
--- a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h
+++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from et_vk_sdpa_qk.wgsl - DO NOT EDIT.
-// wgsl-sha256: cdaa33d5b652395821e1d74222b113e72391b9d618358199abd0772cb821ef40
+// wgsl-sha256: b39754ade084ba2ea2994fe74f9add82759de8a3184df3bb0fd00754c21cdaed
 inline constexpr const char* kEtVkSdpaQkWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> attn: array<f32>;
 @group(0) @binding(1) var<storage, read> q: array<vec4<f32>>;
@@ -22,52 +22,137 @@ inline constexpr const char* kEtVkSdpaQkWGSL = R"(
 
 struct Params {
   B: u32,
-  H: u32,
+  Hq: u32,
+  Hkv: u32,
   S_q: u32,
   S_kv: u32,
   D: u32,
+  g: u32,
   has_mask: u32,
+  mask_mode: u32,
+  tensor_layout: u32,
   _pad0: u32,
   scale: f32,
+  bh_lo: u32,
+  bh_count: u32,
+  _pad1: u32,
+  _pad2: u32,
 }
 @group(0) @binding(4) var<uniform> params: Params;
 
 override wg_size: u32 = 64;
 
-// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D],
-// k [B, H, S_kv, D]; q/k viewed as vec4<f32> over D (caller guarantees
-// D % 4 == 0 for every model in scope). Supports asymmetric seq (S_q != S_kv,
-// e.g. Hiera pooled query); when S_q == S_kv this reduces to plain
-// self-attention. ONE thread per (b, h, s) ROW of attn_weights
-// [B, H, S_q, S_kv]; the thread loops over c (0..S_kv) and d4 (0..D/4). Row
-// count = B*H*S_q stays well under the 65535 1D dispatch limit for any ViT.
+// TMxTN register tile; mirrors sdpa/sdpa_compute_attn_weights.wgsl (not a knob).
+const TM: u32 = 8u;
+const TN: u32 = 4u;
+
+fn q_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hq + h) * params.S_q + s) * d4_count;
+  }
+  return ((b * params.S_q + s) * params.Hq + h) * d4_count;
+}
+
+fn k_row4(b: u32, h: u32, s: u32, d4_count: u32) -> u32 {
+  if (params.tensor_layout == 0u) {
+    return ((b * params.Hkv + h) * params.S_kv + s) * d4_count;
+  }
+  return ((b * params.S_kv + s) * params.Hkv + h) * d4_count;
+}
+
+fn load_q4(b: u32, h: u32, s: u32, d4: u32, d4_count: u32) -> vec4<f32> {
+  if (s >= params.S_q) {
+    return vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+  return q[q_row4(b, h, s, d4_count) + d4];
+}
+
+fn load_k4(b: u32, h: u32, c: u32, d4: u32, d4_count: u32) -> vec4<f32> {
+  if (c >= params.S_kv) {
+    return vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+  return k[k_row4(b, h, c, d4_count) + d4];
+}
+
+fn store_qk(
+    scratch_row: u32,
+    absolute_row: u32,
+    s: u32,
+    c: u32,
+    raw: f32) {
+  if (c >= params.S_kv) {
+    return;
+  }
+  var val = raw * params.scale;
+  if (params.has_mask != 0u) {
+    if (params.mask_mode == 1u) {
+      val = val + mask[params.S_kv * s + c];
+    } else {
+      val = val + mask[absolute_row + c];
+    }
+  }
+  attn[scratch_row + c] = val;
+}
+
 @compute @workgroup_size(wg_size)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let num_rows = params.B * params.H * params.S_q;
-  let row = gid.x;
-  if (row >= num_rows) {
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  let nrt = (params.S_q + TM - 1u) / TM;
+  let nct = (params.S_kv + TN - 1u) / TN;
+  let tiles = nrt * nct;
+  let total = tiles * params.bh_count;
+  // 2D dispatch fold: recover the linear tile index across x/y.
+  let idx = gid.x + gid.y * (nwg.x * wg_size);
+  if (idx >= total) {
     return;
   }
-  let s = row % params.S_q;
-  let h = (row / params.S_q) % params.H;
-  let b = row / (params.S_q * params.H);
 
+  // Tile within one (b, h) so a tile never straddles a head boundary.
+  let local_bh = idx / tiles;
+  let bh = params.bh_lo + local_bh;
+  let rem = idx % tiles;
+  let h = bh % params.Hq;
+  let b = bh / params.Hq;
+  let s0 = (rem / nct) * TM;
+  let c0 = (rem % nct) * TN;
+  let kv_h = h / params.g;
   let d4_count = params.D / 4u;
-  let qbase4 = ((b * params.H + h) * params.S_q + s) * d4_count;
-  let kblock4 = (b * params.H + h) * params.S_kv; // first K row of this (b, h)
-  let arow = ((b * params.H + h) * params.S_q + s) * params.S_kv;
-
-  for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) {
-    let kbase4 = (kblock4 + c) * d4_count;
-    var acc: f32 = 0.0;
-    for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
-      acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]);
+
+  var acc: array<vec4<f32>, TM>;
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    acc[i] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+
+  for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) {
+    var qv: array<vec4<f32>, TM>;
+    var kv: array<vec4<f32>, TN>;
+    for (var i: u32 = 0u; i < TM; i = i + 1u) {
+      qv[i] = load_q4(b, h, s0 + i, d4, d4_count);
+    }
+    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+      kv[j] = load_k4(b, kv_h, c0 + j, d4, d4_count);
     }
-    acc = acc * params.scale;
-    if (params.has_mask != 0u) {
-      acc = acc + mask[arow + c];
+    for (var i: u32 = 0u; i < TM; i = i + 1u) {
+      acc[i] = acc[i] +
+          vec4<f32>(
+              dot(qv[i], kv[0]),
+              dot(qv[i], kv[1]),
+              dot(qv[i], kv[2]),
+              dot(qv[i], kv[3]));
+    }
+  }
+
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    let s = s0 + i;
+    if (s < params.S_q) {
+      let absolute_row = (bh * params.S_q + s) * params.S_kv;
+      let scratch_row = (local_bh * params.S_q + s) * params.S_kv;
+      let av = acc[i];
+      store_qk(scratch_row, absolute_row, s, c0 + 0u, av.x);
+      store_qk(scratch_row, absolute_row, s, c0 + 1u, av.y);
+      store_qk(scratch_row, absolute_row, s, c0 + 2u, av.z);
+      store_qk(scratch_row, absolute_row, s, c0 + 3u, av.w);
     }
-    attn[arow + c] = acc;
   }
 }
 )";
diff --git a/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp
index 3b113b85fc2..89b9415e55a 100644
--- a/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp
+++ b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp
@@ -9,10 +9,13 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
-#include <executorch/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row_wgsl.h>
 
 #include <webgpu/webgpu.h>
 
+#include <algorithm>
 #include <cstdint>
 #include <stdexcept>
 #include <vector>
@@ -21,198 +24,231 @@ namespace executorch::backends::webgpu {
 
 namespace {
 
-struct Dq8caParams {
-  uint32_t M;
-  uint32_t N;
-  uint32_t K;
-  uint32_t K_packed;
-  uint32_t group_size;
-  uint32_t padded_N;
-  uint32_t has_bias;
+struct QuantizeDequantizeParams {
+  uint32_t num_elements;
+  uint32_t num_rows;
+  uint32_t row_width;
   uint32_t _pad;
 };
-static_assert(sizeof(Dq8caParams) == 32, "Dq8caParams must be 32 bytes");
+static_assert(sizeof(QuantizeDequantizeParams) == 16);
+static_assert(
+    kChooseQparamsDq8caFusedWorkgroupSizeX == utils::kCqpQdqFusedInvocations);
 
-constexpr int64_t kTileM = 4; // MUST match TM in linear_dq8ca_q4gsw.wgsl
-constexpr int64_t kTileN = 4; // MUST match TN
-
-// et_vk.linear_dq8ca_q4gsw args (mirrors Vulkan QuantizedLinear.cpp:760):
-// [in, input_scale, input_zp, weight, weight_sums, weight_scales, group_size,
-//  bias, out]. Dynamic per-row int8 activation quant x 4-bit-group symmetric
-// weight. weight_sums (arg 4) is a perf shortcut; this v1 recomputes the sum
-// inline so it is intentionally unused. Static-shape only (no resize hook yet).
-void linear_dq8ca_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
-  const int in_id = args.at(0);
-  const int input_scale_id = args.at(1);
-  const int input_zp_id = args.at(2);
-  const int weight_id = args.at(3);
-  const int scales_id = args.at(5);
-  const int group_size_id = args.at(6);
-  const int bias_id = args.at(7);
-  const int out_id = args.at(8);
-
-  WGPUDevice device = graph.device();
-  const auto& in = graph.get_tensor(in_id);
-  const auto& input_scale = graph.get_tensor(input_scale_id);
-  const auto& input_zp = graph.get_tensor(input_zp_id);
-  const auto& weight = graph.get_tensor(weight_id);
-  const auto& scales = graph.get_tensor(scales_id);
-  const auto& out = graph.get_tensor(out_id);
-
-  if (in.dims.empty() || weight.dims.size() < 2 || scales.dims.size() < 2) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: malformed dims");
-  }
-  if (in.buffer == nullptr || input_scale.buffer == nullptr ||
-      input_zp.buffer == nullptr || weight.buffer == nullptr ||
-      scales.buffer == nullptr || out.buffer == nullptr) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: null buffer binding");
-  }
+struct QuantizeDequantizeState {
+  QuantizeDequantizeParams params;
+  utils::WgCount grid;
+};
 
-  const uint32_t K = static_cast<uint32_t>(in.dims.back());
-  if (K == 0) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: K == 0");
-  }
-  uint64_t in_numel = 1;
-  for (int64_t d : in.dims) {
-    in_numel *= static_cast<uint64_t>(d);
-  }
-  if (in_numel % K != 0) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: input numel % K != 0");
-  }
-  const uint32_t M = static_cast<uint32_t>(in_numel / K);
-  const uint32_t N = static_cast<uint32_t>(weight.dims[0]);
-  const uint32_t K_packed = static_cast<uint32_t>(weight.dims[1]);
-  const uint32_t num_groups = static_cast<uint32_t>(scales.dims[0]);
-  const uint32_t padded_N = static_cast<uint32_t>(scales.dims[1]);
-  if (M == 0 || N == 0) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: M or N == 0");
+QuantizeDequantizeState make_quantize_dequantize_state(
+    WGPUDevice device,
+    const std::vector<int64_t>& input_dims,
+    uint32_t max_rows,
+    uint32_t row_width,
+    uint32_t workgroup_size) {
+  if (input_dims.empty() || input_dims.back() != row_width) {
+    throw std::runtime_error(
+        "WebGPU linear_dq8ca_q4gsw: live row width mismatch");
   }
-  if (K_packed != (K + 1) / 2) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: K_packed must be ceil(K/2)");
+  const uint64_t numel = utils::numel_of(input_dims);
+  if (numel == 0u || numel % row_width != 0u || numel > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU linear_dq8ca_q4gsw: invalid live input numel");
   }
-  if ((static_cast<uint64_t>(N) * K_packed) % 4u != 0u) {
+  const uint64_t rows = numel / row_width;
+  if (rows == 0u || rows > max_rows) {
     throw std::runtime_error(
-        "linear_dq8ca_q4gsw: N*K_packed must be a multiple of 4 (u32-packed)");
+        "WebGPU linear_dq8ca_q4gsw: live rows exceed the build-time max");
   }
 
-  int64_t group_size = 0;
-  if (graph.get_value_type(group_size_id) == WebGPUGraph::ValueType::Int) {
-    group_size = graph.get_int(group_size_id);
-  }
-  if (group_size <= 0) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: group_size <= 0");
+  QuantizeDequantizeState state = {};
+  state.params = {
+      static_cast<uint32_t>(numel), static_cast<uint32_t>(rows), row_width, 0u};
+  state.grid = utils::compute_2d_workgroup_count(
+      device,
+      state.params.num_elements,
+      workgroup_size,
+      "linear_dq8ca_q4gsw_qdq");
+  return state;
+}
+
+utils::WgCount make_cqp_fused_grid(
+    WGPUDevice device,
+    const QuantizeDequantizeParams& params) {
+  const uint32_t grid_x = utils::clamp_workgroup_count(device, params.num_rows);
+  if (grid_x == 0u) {
+    throw std::runtime_error("WebGPU linear_dq8ca_q4gsw(fused): zero dispatch");
   }
-  const uint32_t gs = static_cast<uint32_t>(group_size);
-
-  // fp32-only byte guards; per-row scale (f32[M]) + zp (int8[M]).
-  if (in.nbytes != in_numel * sizeof(float) ||
-      out.nbytes != static_cast<uint64_t>(M) * N * sizeof(float) ||
-      scales.nbytes !=
-          static_cast<uint64_t>(num_groups) * padded_N * sizeof(float) ||
-      weight.nbytes != static_cast<uint64_t>(N) * K_packed) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: fp32/byte-size mismatch");
+  return {grid_x, 1u};
+}
+
+WebGPUGraph::CqpFusionSite claim_cqp_fusion_site(
+    WebGPUGraph& graph,
+    int input_id,
+    int input_scales_id,
+    int input_zero_points_id,
+    uint32_t rows,
+    uint32_t row_width,
+    uint64_t numel) {
+  WebGPUGraph::CqpFusionSite site = graph.claim_cqp_fusion_site(
+      input_id, input_scales_id, input_zero_points_id, rows, row_width);
+  if (!site.valid) {
+    return site;
   }
-  // Per-row activation scale (fp32[M]) + zp (int8[M], packed 4/word in-shader).
-  if (input_scale.nbytes != static_cast<uint64_t>(M) * sizeof(float) ||
-      !input_zp.is_int8 || input_zp.nbytes != static_cast<uint64_t>(M)) {
-    throw std::runtime_error(
-        "linear_dq8ca_q4gsw: input scale fp32[M] / zp int8[M] required");
+  WGPULimits limits = {};
+  if (wgpuDeviceGetLimits(graph.device(), &limits) != WGPUStatus_Success ||
+      !utils::is_cqp_qdq_fusion_eligible(
+          rows,
+          row_width,
+          numel,
+          site.quant_min,
+          site.quant_max,
+          true,
+          true,
+          false,
+          limits.maxComputeInvocationsPerWorkgroup,
+          limits.maxComputeWorkgroupSizeX,
+          limits.maxComputeWorkgroupStorageSize)) {
+    return WebGPUGraph::CqpFusionSite{};
   }
-  // int8 zp is bound word-aligned over a max(nbytes,4) buffer; M in {5,6,7,...}
-  // would bind past the buffer. Mirrors the choose_qparams_affine producer
-  // guard.
-  if (M > 4u && M % 4u != 0u) {
-    throw std::runtime_error(
-        "linear_dq8ca_q4gsw: num_rows must be <=4 or a multiple of 4");
+  return site;
+}
+
+void linear_dq8ca_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int input_id = args.at(0);
+  const int input_scales_id = args.at(1);
+  const int input_zero_points_id = args.at(2);
+  const int weight_id = args.at(3);
+  const int weight_scales_id = args.at(5);
+  const int group_size_id = args.at(6);
+  const int bias_id = args.at(7);
+  const int output_id = args.at(8);
+
+  const auto& input = graph.get_tensor(input_id);
+  const auto& input_scales = graph.get_tensor(input_scales_id);
+  const auto& input_zero_points = graph.get_tensor(input_zero_points_id);
+  if (input.buffer == nullptr || input.dims.empty() || input.is_int ||
+      !utils::is_fp32_tensor(input)) {
+    throw std::runtime_error("WebGPU linear_dq8ca_q4gsw: expected fp32 input");
   }
-  if (num_groups < (K + gs - 1u) / gs || padded_N < N) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: scales dims too small");
+  if (input.dims.back() <= 0 ||
+      static_cast<uint64_t>(input.dims.back()) > UINT32_MAX) {
+    throw std::runtime_error("WebGPU linear_dq8ca_q4gsw: invalid row width");
   }
-
-  uint32_t has_bias = 0;
-  WGPUBuffer bias_buffer = nullptr;
-  uint64_t bias_size = 4;
-  if (graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor) {
-    const auto& bias = graph.get_tensor(bias_id);
-    if (bias.buffer != nullptr && bias.nbytes >= N * sizeof(float)) {
-      has_bias = 1;
-      bias_buffer = bias.buffer;
-      bias_size = bias.nbytes;
-    }
+  const uint32_t row_width = static_cast<uint32_t>(input.dims.back());
+  const uint64_t input_numel = utils::numel_of(input.dims);
+  if (row_width == 0u || input_numel % row_width != 0u) {
+    throw std::runtime_error("WebGPU linear_dq8ca_q4gsw: invalid input shape");
   }
-  if (bias_buffer == nullptr) {
-    bias_buffer = graph.create_scratch_buffer(4);
+  const uint64_t max_rows64 = input_numel / row_width;
+  if (max_rows64 == 0u || max_rows64 > UINT32_MAX) {
+    throw std::runtime_error("WebGPU linear_dq8ca_q4gsw: rows out of range");
   }
-
-  Dq8caParams params = {};
-  params.M = M;
-  params.N = N;
-  params.K = K;
-  params.K_packed = K_packed;
-  params.group_size = gs;
-  params.padded_N = padded_N;
-  params.has_bias = has_bias;
-
-  const int64_t total_tiles =
-      utils::div_up<int64_t>(M, kTileM) * utils::div_up<int64_t>(N, kTileN);
-  if (total_tiles > static_cast<int64_t>(UINT32_MAX)) {
-    throw std::runtime_error("linear_dq8ca_q4gsw: tile count exceeds u32");
+  const uint32_t max_rows = static_cast<uint32_t>(max_rows64);
+  if (input_scales.buffer == nullptr || input_zero_points.buffer == nullptr ||
+      input_scales.is_int || input_scales.elem_size != sizeof(float) ||
+      !input_zero_points.is_int8 ||
+      input_zero_points.elem_size != sizeof(int8_t) ||
+      input_scales.dims != input_zero_points.dims ||
+      utils::numel_of(input_scales.dims) != max_rows ||
+      input_scales.nbytes != max_rows * sizeof(float) ||
+      input_zero_points.nbytes != max_rows) {
+    throw std::runtime_error(
+        "WebGPU linear_dq8ca_q4gsw: invalid per-row qparams");
   }
-  const uint32_t wg_size =
-      utils::clamp_workgroup_size(device, kLinearDq8caQ4gswWorkgroupSizeX);
-  const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
-      device,
-      static_cast<uint32_t>(total_tiles),
-      wg_size,
-      "linear_dq8ca_q4gsw");
-
-  WGPUConstantEntry wg_size_constant = {};
-  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
-  wg_size_constant.value = static_cast<double>(wg_size);
 
-  WGPUBuffer params_buf =
-      utils::make_uniform(device, &params, sizeof(Dq8caParams));
-  graph.add_uniform_buffer_bytes(sizeof(Dq8caParams));
+  WebGPUGraph::ScopedScratch scratch(
+      &graph, graph.acquire_scratch(input.nbytes));
+  const WebGPUGraph::CqpFusionSite fusion = claim_cqp_fusion_site(
+      graph,
+      input_id,
+      input_scales_id,
+      input_zero_points_id,
+      max_rows,
+      row_width,
+      input_numel);
+  const bool use_fused = fusion.valid;
+  const uint32_t workgroup_size = utils::clamp_workgroup_size(
+      graph.device(), kQuantizeDequantizePerRowWorkgroupSizeX);
+  const QuantizeDequantizeState initial_state = make_quantize_dequantize_state(
+      graph.device(), input.dims, max_rows, row_width, workgroup_size);
+  const utils::WgCount initial_grid = use_fused
+      ? make_cqp_fused_grid(graph.device(), initial_state.params)
+      : initial_state.grid;
+  WGPUBuffer uniform = graph.make_uniform_buffer(
+      &initial_state.params, sizeof(QuantizeDequantizeParams));
 
-  // 0 out(rw), 1 in, 2 input_scale, 3 input_zp, 4 weight, 5 scales, 6 bias
-  // (ro), 7 uniform.
+  WGPUConstantEntry workgroup_constant = {};
+  workgroup_constant.key = {"wg_size", WGPU_STRLEN};
+  workgroup_constant.value = static_cast<double>(workgroup_size);
   utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
-      device,
-      kLinearDq8caQ4gswWGSL,
+      graph.device(),
+      use_fused ? kChooseQparamsDq8caFusedWGSL : kQuantizeDequantizePerRowWGSL,
       {
-          {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
-          {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes},
-          {2,
+          {0, WGPUBufferBindingType_Storage, scratch.buf, input.nbytes},
+          {1,
            WGPUBufferBindingType_ReadOnlyStorage,
-           input_scale.buffer,
-           input_scale.nbytes},
-          // int8 zp bound as array<u32>; round to a multiple of 4 (buffer is
-          // >=4 bytes).
+           input.buffer,
+           input.nbytes},
+          {2,
+           use_fused ? WGPUBufferBindingType_Storage
+                     : WGPUBufferBindingType_ReadOnlyStorage,
+           input_scales.buffer,
+           input_scales.nbytes},
           {3,
-           WGPUBufferBindingType_ReadOnlyStorage,
-           input_zp.buffer,
-           ((input_zp.nbytes + 3u) / 4u) * 4u},
+           use_fused ? WGPUBufferBindingType_Storage
+                     : WGPUBufferBindingType_ReadOnlyStorage,
+           input_zero_points.buffer,
+           std::max(((input_zero_points.nbytes + 3u) / 4u) * 4u, size_t(4))},
           {4,
-           WGPUBufferBindingType_ReadOnlyStorage,
-           weight.buffer,
-           weight.nbytes},
-          {5,
-           WGPUBufferBindingType_ReadOnlyStorage,
-           scales.buffer,
-           scales.nbytes},
-          {6, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size},
-          {7, WGPUBufferBindingType_Uniform, params_buf, sizeof(Dq8caParams)},
+           WGPUBufferBindingType_Uniform,
+           uniform,
+           sizeof(QuantizeDequantizeParams)},
       },
-      &wg_size_constant,
-      1);
+      use_fused ? nullptr : &workgroup_constant,
+      use_fused ? 0u : 1u);
 
-  graph.add_dispatch(
+  const size_t dispatch_index = graph.add_dispatch(
       {bundle.pipeline,
        bundle.bind_group,
-       workgroup_count.x,
-       "linear_dq8ca_q4gsw",
-       workgroup_count.y});
-  graph.own_uniform_buffer(params_buf);
+       initial_grid.x,
+       use_fused ? "linear_dq8ca_q4gsw_cqp_fused" : "linear_dq8ca_q4gsw_qdq",
+       initial_grid.y});
+  graph.add_tensor_resize_hook(
+      input_id,
+      [input_id,
+       max_rows,
+       row_width,
+       workgroup_size,
+       use_fused,
+       dispatch_index,
+       uniform](WebGPUGraph& g) {
+        const QuantizeDequantizeState state = make_quantize_dequantize_state(
+            g.device(),
+            g.cur_dims(input_id),
+            max_rows,
+            row_width,
+            workgroup_size);
+        wgpuQueueWriteBuffer(
+            g.queue(), uniform, 0, &state.params, sizeof(state.params));
+        const utils::WgCount grid = use_fused
+            ? make_cqp_fused_grid(g.device(), state.params)
+            : state.grid;
+        auto& dispatch = g.dispatch_at(dispatch_index);
+        dispatch.workgroup_count_x = grid.x;
+        dispatch.workgroup_count_y = grid.y;
+      });
+  if (use_fused) {
+    auto& producer = graph.dispatch_at(fusion.dispatch_index);
+    producer.workgroup_count_x = 0u;
+    producer.workgroup_count_y = 0u;
+    *fusion.producer_elided = true;
+  }
+
+  graph.own_uniform_buffer(uniform);
+  const std::vector<int> q4_args = {
+      input_id, weight_id, weight_scales_id, group_size_id, bias_id, output_id};
+  q4gsw_linear_impl_with_input_buffer(
+      graph, q4_args, scratch.buf, input.nbytes);
 }
 
 } // namespace
diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp
index 27bff334dd3..9637ea05391 100644
--- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp
+++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp
@@ -11,6 +11,7 @@
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/TensorMeta.h>
 #include <executorch/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.h>
 
 #include <webgpu/webgpu.h>
 
@@ -28,6 +29,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const int in2_id = args.at(1);
   const int out_id = args.at(2);
 
+  if (fusion::try_fuse_scale(graph, in1_id, in2_id, out_id)) {
+    return;
+  }
+
   WGPUDevice device = graph.device();
 
   const auto& in1_tensor = graph.get_tensor(in1_id);
diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
index fe60e8377a1..dcca24b357d 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
+++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
@@ -11,6 +11,7 @@
 #include <executorch/backends/webgpu/runtime/WebGPUShaderRegistry.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.h>
 
 #include <webgpu/webgpu.h>
 
@@ -351,7 +352,11 @@ void resize_q4gsw(WebGPUGraph& graph, const Q4gswResizeContext& context) {
 }
 
 // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out].
-void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+void q4gsw_linear_impl_with_input_buffer_internal(
+    WebGPUGraph& graph,
+    const std::vector<int>& args,
+    WGPUBuffer input_buffer,
+    uint64_t input_nbytes) {
   const int in_id = args.at(0);
   const int weight_id = args.at(1);
   const int scales_id = args.at(2);
@@ -406,7 +411,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       static_cast<uint64_t>(num_groups) * static_cast<uint64_t>(padded_N);
   const uint64_t weight_numel =
       static_cast<uint64_t>(N) * static_cast<uint64_t>(K_packed);
-  if (in.nbytes != in_numel * sizeof(float) ||
+  if (input_buffer == nullptr || input_nbytes != in_numel * sizeof(float) ||
       out.nbytes != static_cast<uint64_t>(M) * N * sizeof(float) ||
       scales.nbytes != scales_numel * sizeof(float) ||
       weight.nbytes != weight_numel) {
@@ -538,7 +543,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   WGPUBuffer params_buffer = graph.create_params_buffer(initial_state.params);
   const std::vector<utils::BindingSpec> bindings = {
       {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
-      {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes},
+      {1, WGPUBufferBindingType_ReadOnlyStorage, input_buffer, input_nbytes},
       {2, WGPUBufferBindingType_ReadOnlyStorage, weight.buffer, weight.nbytes},
       {3, WGPUBufferBindingType_ReadOnlyStorage, scales.buffer, scales.nbytes},
       {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size},
@@ -665,6 +670,20 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
 } // namespace
 
+void q4gsw_linear_impl_with_input_buffer(
+    WebGPUGraph& graph,
+    const std::vector<int>& args,
+    WGPUBuffer input_buffer,
+    uint64_t input_nbytes) {
+  q4gsw_linear_impl_with_input_buffer_internal(
+      graph, args, input_buffer, input_nbytes);
+}
+
+void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const auto& input = graph.get_tensor(args.at(0));
+  q4gsw_linear_impl_with_input_buffer(graph, args, input.buffer, input.nbytes);
+}
+
 WEBGPU_REGISTER_OPERATORS {
   WEBGPU_REGISTER_OP(et_vk.linear_q4gsw.default, q4gsw_linear_impl);
 }
diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.h b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.h
new file mode 100644
index 00000000000..f7e6efff7da
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.h
@@ -0,0 +1,26 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+class WebGPUGraph;
+
+void q4gsw_linear_impl_with_input_buffer(
+    WebGPUGraph& graph,
+    const std::vector<int>& args,
+    WGPUBuffer input_buffer,
+    uint64_t input_nbytes);
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused.wgsl b/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused.wgsl
new file mode 100644
index 00000000000..fdbb033fce7
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused.wgsl
@@ -0,0 +1,164 @@
+// 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.
+
+struct Params {
+  num_elements: u32,
+  num_rows: u32,
+  row_width: u32,
+  _pad: u32,
+};
+
+@group(0) @binding(0) var<storage, read_write> output: array<f32>;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read_write> scales: array<f32>;
+@group(0) @binding(3) var<storage, read_write> zero_points: array<atomic<u32>>;
+@group(0) @binding(4) var<uniform> params: Params;
+
+const WG: u32 = 256u;
+const QUANT_MIN: i32 = -128;
+const QUANT_MAX: i32 = 127;
+const SMALL_SCALE_THRESHOLD: f32 = 6.1e-5;
+
+var<workgroup> min_values: array<f32, WG>;
+var<workgroup> max_values: array<f32, WG>;
+
+fn reciprocal_is_infinite(value: f32) -> bool {
+  // WGSL has no portable isinf builtin. This exponent/mantissa check is the
+  // exact f32 equivalent used for Vulkan's isinf(1.0 / scale) condition.
+  return (bitcast<u32>(1.0 / value) & 0x7fffffffu) == 0x7f800000u;
+}
+
+fn calculate_scale_and_zero_point(
+    input_min: f32,
+    input_max: f32) -> vec2<f32> {
+  var min_value = min(input_min, 0.0);
+  var max_value = max(input_max, 0.0);
+  let qmin = f32(QUANT_MIN);
+  let qmax = f32(QUANT_MAX);
+  var scale = (max_value - min_value) / (qmax - qmin);
+  if (scale == 0.0 || reciprocal_is_infinite(scale)) {
+    scale = 0.1;
+  }
+  if (scale < SMALL_SCALE_THRESHOLD) {
+    let original_scale = scale;
+    scale = SMALL_SCALE_THRESHOLD;
+    if (min_value == 0.0) {
+      max_value = SMALL_SCALE_THRESHOLD * (qmax - qmin);
+    } else if (max_value == 0.0) {
+      min_value = -SMALL_SCALE_THRESHOLD * (qmax - qmin);
+    } else {
+      let amplifier = SMALL_SCALE_THRESHOLD / original_scale;
+      min_value *= amplifier;
+      max_value *= amplifier;
+    }
+  }
+
+  let zero_point_from_min = qmin - min_value / scale;
+  let zero_point_from_max = qmax - max_value / scale;
+  let zero_point_from_min_error = abs(qmin) - abs(min_value / scale);
+  let zero_point_from_max_error = abs(qmax) - abs(max_value / scale);
+  var initial_zero_point = zero_point_from_max;
+  if (zero_point_from_min_error < zero_point_from_max_error) {
+    initial_zero_point = zero_point_from_min;
+  }
+  var nudged_zero_point: i32;
+  if (initial_zero_point < qmin) {
+    nudged_zero_point = QUANT_MIN;
+  } else if (initial_zero_point > qmax) {
+    nudged_zero_point = QUANT_MAX;
+  } else {
+    nudged_zero_point = i32(round(initial_zero_point));
+  }
+  return vec2<f32>(scale, f32(nudged_zero_point));
+}
+
+fn load_zero_point(row: u32) -> i32 {
+  let word = atomicLoad(&zero_points[row / 4u]);
+  let byte = (word >> ((row % 4u) * 8u)) & 0xffu;
+  return select(i32(byte), i32(byte) - 256, byte >= 128u);
+}
+
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) ngrp: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  let lane = lid.x;
+  var row = wid.x;
+  loop {
+    if (row >= params.num_rows) {
+      break;
+    }
+    let row_start = row * params.row_width;
+
+    var local_min = input[row_start];
+    var local_max = input[row_start];
+    var col = lane;
+    while (col < params.row_width) {
+      let value = input[row_start + col];
+      local_min = min(local_min, value);
+      local_max = max(local_max, value);
+      col += WG;
+    }
+
+    min_values[lane] = local_min;
+    max_values[lane] = local_max;
+    workgroupBarrier();
+
+    var stride = WG / 2u;
+    while (stride > 0u) {
+      if (lane < stride) {
+        min_values[lane] = min(min_values[lane], min_values[lane + stride]);
+        max_values[lane] = max(max_values[lane], max_values[lane + stride]);
+      }
+      workgroupBarrier();
+      stride /= 2u;
+    }
+
+    if (lane == 0u) {
+      let qparams =
+          calculate_scale_and_zero_point(min_values[0], max_values[0]);
+      scales[row] = qparams.x;
+      let zero_point_byte = u32(i32(qparams.y)) & 0xffu;
+      let pack = row / 4u;
+      let shift = (row % 4u) * 8u;
+      var clear_mask = ~(0xffu << shift);
+      if (row % 4u == 0u) {
+        for (var tail = 1u; tail < 4u; tail++) {
+          if (row + tail >= params.num_rows) {
+            clear_mask &= ~(0xffu << (tail * 8u));
+          }
+        }
+      }
+      atomicAnd(&zero_points[pack], clear_mask);
+      atomicOr(&zero_points[pack], zero_point_byte << shift);
+    }
+
+    storageBarrier();
+    workgroupBarrier();
+
+    let scale = scales[row];
+    let zero_point = load_zero_point(row);
+    var elem = lane;
+    while (elem < params.row_width) {
+      let index = row_start + elem;
+      let quantized = clamp(
+          round(input[index] * (1.0 / scale)) + f32(zero_point),
+          -128.0,
+          127.0);
+      output[index] = (quantized - f32(zero_point)) * scale;
+      elem += WG;
+    }
+
+    workgroupBarrier();
+    row += ngrp.x;
+  }
+}
+// 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.
diff --git a/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused_wgsl.h
new file mode 100644
index 00000000000..a431964dbd5
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/choose_qparams_dq8ca_fused_wgsl.h
@@ -0,0 +1,188 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from choose_qparams_dq8ca_fused.wgsl - DO NOT EDIT.
+// wgsl-sha256: ce8bd3a728dae5657827e81fb1bdc9eb60b3d7f868a9e1f01412e8704c240a31
+inline constexpr const char* kChooseQparamsDq8caFusedWGSL = R"(
+// 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.
+
+struct Params {
+  num_elements: u32,
+  num_rows: u32,
+  row_width: u32,
+  _pad: u32,
+};
+
+@group(0) @binding(0) var<storage, read_write> output: array<f32>;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read_write> scales: array<f32>;
+@group(0) @binding(3) var<storage, read_write> zero_points: array<atomic<u32>>;
+@group(0) @binding(4) var<uniform> params: Params;
+
+const WG: u32 = 256u;
+const QUANT_MIN: i32 = -128;
+const QUANT_MAX: i32 = 127;
+const SMALL_SCALE_THRESHOLD: f32 = 6.1e-5;
+
+var<workgroup> min_values: array<f32, WG>;
+var<workgroup> max_values: array<f32, WG>;
+
+fn reciprocal_is_infinite(value: f32) -> bool {
+  // WGSL has no portable isinf builtin. This exponent/mantissa check is the
+  // exact f32 equivalent used for Vulkan's isinf(1.0 / scale) condition.
+  return (bitcast<u32>(1.0 / value) & 0x7fffffffu) == 0x7f800000u;
+}
+
+fn calculate_scale_and_zero_point(
+    input_min: f32,
+    input_max: f32) -> vec2<f32> {
+  var min_value = min(input_min, 0.0);
+  var max_value = max(input_max, 0.0);
+  let qmin = f32(QUANT_MIN);
+  let qmax = f32(QUANT_MAX);
+  var scale = (max_value - min_value) / (qmax - qmin);
+  if (scale == 0.0 || reciprocal_is_infinite(scale)) {
+    scale = 0.1;
+  }
+  if (scale < SMALL_SCALE_THRESHOLD) {
+    let original_scale = scale;
+    scale = SMALL_SCALE_THRESHOLD;
+    if (min_value == 0.0) {
+      max_value = SMALL_SCALE_THRESHOLD * (qmax - qmin);
+    } else if (max_value == 0.0) {
+      min_value = -SMALL_SCALE_THRESHOLD * (qmax - qmin);
+    } else {
+      let amplifier = SMALL_SCALE_THRESHOLD / original_scale;
+      min_value *= amplifier;
+      max_value *= amplifier;
+    }
+  }
+
+  let zero_point_from_min = qmin - min_value / scale;
+  let zero_point_from_max = qmax - max_value / scale;
+  let zero_point_from_min_error = abs(qmin) - abs(min_value / scale);
+  let zero_point_from_max_error = abs(qmax) - abs(max_value / scale);
+  var initial_zero_point = zero_point_from_max;
+  if (zero_point_from_min_error < zero_point_from_max_error) {
+    initial_zero_point = zero_point_from_min;
+  }
+  var nudged_zero_point: i32;
+  if (initial_zero_point < qmin) {
+    nudged_zero_point = QUANT_MIN;
+  } else if (initial_zero_point > qmax) {
+    nudged_zero_point = QUANT_MAX;
+  } else {
+    nudged_zero_point = i32(round(initial_zero_point));
+  }
+  return vec2<f32>(scale, f32(nudged_zero_point));
+}
+
+fn load_zero_point(row: u32) -> i32 {
+  let word = atomicLoad(&zero_points[row / 4u]);
+  let byte = (word >> ((row % 4u) * 8u)) & 0xffu;
+  return select(i32(byte), i32(byte) - 256, byte >= 128u);
+}
+
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) ngrp: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  let lane = lid.x;
+  var row = wid.x;
+  loop {
+    if (row >= params.num_rows) {
+      break;
+    }
+    let row_start = row * params.row_width;
+
+    var local_min = input[row_start];
+    var local_max = input[row_start];
+    var col = lane;
+    while (col < params.row_width) {
+      let value = input[row_start + col];
+      local_min = min(local_min, value);
+      local_max = max(local_max, value);
+      col += WG;
+    }
+
+    min_values[lane] = local_min;
+    max_values[lane] = local_max;
+    workgroupBarrier();
+
+    var stride = WG / 2u;
+    while (stride > 0u) {
+      if (lane < stride) {
+        min_values[lane] = min(min_values[lane], min_values[lane + stride]);
+        max_values[lane] = max(max_values[lane], max_values[lane + stride]);
+      }
+      workgroupBarrier();
+      stride /= 2u;
+    }
+
+    if (lane == 0u) {
+      let qparams =
+          calculate_scale_and_zero_point(min_values[0], max_values[0]);
+      scales[row] = qparams.x;
+      let zero_point_byte = u32(i32(qparams.y)) & 0xffu;
+      let pack = row / 4u;
+      let shift = (row % 4u) * 8u;
+      var clear_mask = ~(0xffu << shift);
+      if (row % 4u == 0u) {
+        for (var tail = 1u; tail < 4u; tail++) {
+          if (row + tail >= params.num_rows) {
+            clear_mask &= ~(0xffu << (tail * 8u));
+          }
+        }
+      }
+      atomicAnd(&zero_points[pack], clear_mask);
+      atomicOr(&zero_points[pack], zero_point_byte << shift);
+    }
+
+    storageBarrier();
+    workgroupBarrier();
+
+    let scale = scales[row];
+    let zero_point = load_zero_point(row);
+    var elem = lane;
+    while (elem < params.row_width) {
+      let index = row_start + elem;
+      let quantized = clamp(
+          round(input[index] * (1.0 / scale)) + f32(zero_point),
+          -128.0,
+          127.0);
+      output[index] = (quantized - f32(zero_point)) * scale;
+      elem += WG;
+    }
+
+    workgroupBarrier();
+    row += ngrp.x;
+  }
+}
+// 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.
+)";
+
+inline constexpr uint32_t kChooseQparamsDq8caFusedWorkgroupSizeX = 256;
+inline constexpr uint32_t kChooseQparamsDq8caFusedWorkgroupSizeY = 1;
+inline constexpr uint32_t kChooseQparamsDq8caFusedWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row.wgsl b/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row.wgsl
new file mode 100644
index 00000000000..6407ff7b1d6
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row.wgsl
@@ -0,0 +1,50 @@
+// 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.
+
+struct Params {
+  num_elements: u32,
+  num_rows: u32,
+  row_width: u32,
+  _pad: u32,
+};
+
+@group(0) @binding(0) var<storage, read_write> output: array<f32>;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read> scales: array<f32>;
+@group(0) @binding(3) var<storage, read> zero_points: array<u32>;
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+fn load_zero_point(row: u32) -> i32 {
+  let word = zero_points[row / 4u];
+  let byte = (word >> ((row % 4u) * 8u)) & 0xffu;
+  return select(i32(byte), i32(byte) - 256, byte >= 128u);
+}
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
+  let index = gid.x + gid.y * nwg.x * wg_size;
+  if (index >= params.num_elements) {
+    return;
+  }
+  let row = index / params.row_width;
+  if (row >= params.num_rows) {
+    return;
+  }
+  let scale = scales[row];
+  let zero_point = load_zero_point(row);
+  let quantized = clamp(
+      round(input[index] * (1.0 / scale)) + f32(zero_point), -128.0, 127.0);
+  output[index] = (quantized - f32(zero_point)) * scale;
+}
+// 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.
diff --git a/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row_wgsl.h
new file mode 100644
index 00000000000..b0fb07d0a82
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/quantize_dequantize_per_row_wgsl.h
@@ -0,0 +1,74 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from quantize_dequantize_per_row.wgsl - DO NOT EDIT.
+// wgsl-sha256: 5ad80c6834f48ea9afd45982b7294fbc9049feda520bebdfaf11932e5464d34a
+inline constexpr const char* kQuantizeDequantizePerRowWGSL = R"(
+// 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.
+
+struct Params {
+  num_elements: u32,
+  num_rows: u32,
+  row_width: u32,
+  _pad: u32,
+};
+
+@group(0) @binding(0) var<storage, read_write> output: array<f32>;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read> scales: array<f32>;
+@group(0) @binding(3) var<storage, read> zero_points: array<u32>;
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+fn load_zero_point(row: u32) -> i32 {
+  let word = zero_points[row / 4u];
+  let byte = (word >> ((row % 4u) * 8u)) & 0xffu;
+  return select(i32(byte), i32(byte) - 256, byte >= 128u);
+}
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
+  let index = gid.x + gid.y * nwg.x * wg_size;
+  if (index >= params.num_elements) {
+    return;
+  }
+  let row = index / params.row_width;
+  if (row >= params.num_rows) {
+    return;
+  }
+  let scale = scales[row];
+  let zero_point = load_zero_point(row);
+  let quantized = clamp(
+      round(input[index] * (1.0 / scale)) + f32(zero_point), -128.0, 127.0);
+  output[index] = (quantized - f32(zero_point)) * scale;
+}
+// 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.
+)";
+
+inline constexpr uint32_t kQuantizeDequantizePerRowWorkgroupSizeX = 64;
+inline constexpr uint32_t kQuantizeDequantizePerRowWorkgroupSizeY = 1;
+inline constexpr uint32_t kQuantizeDequantizePerRowWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
index 50b8828de93..9a4e0b414bc 100644
--- a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
+++ b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
@@ -9,6 +9,7 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.h>
 #include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h>
 
@@ -52,16 +53,15 @@ void resize_rms_norm(
   if (rows == 0) {
     throw std::runtime_error("WebGPU rms_norm: zero rows");
   }
-  if (rows > utils::queried_max_workgroups(g.device())) {
-    throw std::runtime_error(
-        "WebGPU rms_norm: num_rows exceeds the 1D dispatch limit");
-  }
   RmsNormParams p = {};
   p.num_rows = rows;
   p.row_width = row_width;
   p.epsilon = epsilon;
   wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
-  g.dispatch_at(dispatch_idx).workgroup_count_x = rows;
+  const utils::WgCount wg =
+      utils::compute_2d_workgroup_count(g.device(), rows, 1, "rms_norm");
+  g.dispatch_at(dispatch_idx).workgroup_count_x = wg.x;
+  g.dispatch_at(dispatch_idx).workgroup_count_y = wg.y;
   g.set_cur_dims(out_id, d);
 }
 
@@ -84,30 +84,32 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
   // row_width = last dim; num_rows = product of the rest (PyTorch NCHW order)
   const auto& in_tensor = graph.get_tensor(in_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+  const auto& weight_tensor = graph.get_tensor(weight_id);
   if (in_tensor.dims.empty() || in_tensor.nbytes == 0) {
     throw std::runtime_error("WebGPU rms_norm: empty input");
   }
-  const uint32_t row_width = static_cast<uint32_t>(in_tensor.dims.back());
-  if (row_width == 0) {
-    throw std::runtime_error("WebGPU rms_norm: zero row width");
+  if (in_tensor.dims.back() <= 0 ||
+      static_cast<uint64_t>(in_tensor.dims.back()) > UINT32_MAX) {
+    throw std::runtime_error("WebGPU rms_norm: invalid row width");
   }
-  uint64_t in_numel = 1;
-  for (int64_t d : in_tensor.dims) {
-    in_numel *= static_cast<uint64_t>(d);
+  const uint32_t row_width = static_cast<uint32_t>(in_tensor.dims.back());
+  const uint64_t in_numel = utils::numel_of(in_tensor.dims);
+  if (!utils::is_fp32_tensor(in_tensor) || !utils::is_fp32_tensor(out_tensor) ||
+      !utils::is_fp32_tensor(weight_tensor) ||
+      out_tensor.dims != in_tensor.dims ||
+      utils::numel_of(weight_tensor.dims) != row_width) {
+    throw std::runtime_error(
+        "WebGPU rms_norm: expected fp32 input/output and row-width weight");
   }
-  // fp32-only shader: bail if the bytes don't match an fp32 element count.
-  if (in_tensor.nbytes != in_numel * sizeof(float)) {
-    throw std::runtime_error("WebGPU rms_norm: fp32-only (byte-size mismatch)");
+  if (in_numel % row_width != 0u || in_numel / row_width == 0u ||
+      in_numel / row_width > UINT32_MAX) {
+    throw std::runtime_error("WebGPU rms_norm: invalid row count");
   }
   const uint32_t num_rows = static_cast<uint32_t>(in_numel / row_width);
-  if (num_rows == 0) {
-    throw std::runtime_error("WebGPU rms_norm: zero rows");
-  }
-  // Validate the 1D dispatch limit before allocating any GPU objects.
-  if (num_rows > utils::queried_max_workgroups(device)) {
-    throw std::runtime_error(
-        "WebGPU rms_norm: num_rows exceeds the 1D dispatch limit");
-  }
+  // Rows can exceed the per-dim grid cap (QK-norm at prefill), so fold x/y.
+  const utils::WgCount wg_count =
+      utils::compute_2d_workgroup_count(device, num_rows, 1, "rms_norm");
 
   // Create uniform buffer for params
   RmsNormParams params = {};
@@ -144,8 +146,6 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
   wg_size_constant.value = static_cast<double>(wg_size);
 
-  const auto& out_tensor = graph.get_tensor(out_id);
-  const auto& weight_tensor = graph.get_tensor(weight_id);
   utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
       device,
       shader_src,
@@ -177,8 +177,24 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   static_assert(
       kRmsNormVec4WorkgroupSizeX == 64,
       "kRmsNormVec4WorkgroupSizeX must match override wg_size default (64)");
-  const size_t dispatch_idx =
-      graph.add_dispatch({bundle.pipeline, bundle.bind_group, num_rows});
+  const size_t dispatch_idx = graph.add_dispatch(
+      {bundle.pipeline, bundle.bind_group, wg_count.x, "rms_norm", wg_count.y});
+
+  // Offer this dispatch to the add/mul merge; vec4 at 64-wide only.
+  if (use_vec4 && wg_size == kRmsNormVec4WorkgroupSizeX) {
+    fusion::record_rms_norm(
+        graph,
+        in_id,
+        weight_id,
+        out_id,
+        num_rows,
+        row_width,
+        dispatch_idx,
+        uniform_buffer,
+        bundle.bind_group);
+  } else {
+    fusion::invalidate_record(graph);
+  }
 
   // Dynamic shapes: recompute num_rows + rewrite the UBO for the live input.
   WGPUBuffer params_buf = uniform_buffer;
diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.cpp
new file mode 100644
index 00000000000..66693ccddb6
--- /dev/null
+++ b/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.cpp
@@ -0,0 +1,275 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.h>
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_wgsl.h>
+
+#include <stdexcept>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+namespace fusion {
+
+namespace {
+
+using Record = WebGPUGraph::RmsFusionSite;
+
+uint64_t numel_of_tensor(const WebGPUTensor& t) {
+  return utils::numel_of(t.dims);
+}
+
+bool still_last(const WebGPUGraph& graph, const Record& r) {
+  return r.valid && r.dispatch_index + 1u == graph.num_dispatches();
+}
+
+void resize_fused_add_outputs(
+    WebGPUGraph& graph,
+    int rms_in_id,
+    int resid_id,
+    int addout_id,
+    int scale_id = -1,
+    int scaleout_id = -1) {
+  const auto& rms_dims = graph.cur_dims(rms_in_id);
+  RmsResizeCheck check;
+  check.residual_shape_matches = graph.cur_dims(resid_id) == rms_dims;
+  check.has_scale = scale_id >= 0;
+  check.scale_numel =
+      check.has_scale ? utils::numel_of(graph.cur_dims(scale_id)) : 0u;
+  if (const char* reason = rms_resize_reject_reason(check)) {
+    throw std::runtime_error(std::string("WebGPU rms fusion: ") + reason);
+  }
+  graph.set_cur_dims(addout_id, rms_dims);
+  if (scaleout_id >= 0) {
+    graph.set_cur_dims(scaleout_id, rms_dims);
+  }
+}
+
+// Replace dispatches_[idx]'s pipeline + bind group only after the complete
+// replacement bundle has been created successfully.
+void install(
+    WebGPUGraph& graph,
+    size_t dispatch_idx,
+    const char* wgsl,
+    const std::vector<utils::BindingSpec>& bindings) {
+  utils::ComputePipelineBundle replacement =
+      utils::make_compute_pipeline(graph.device(), wgsl, bindings);
+
+  WebGPUDispatch& d = graph.dispatch_at(dispatch_idx);
+  WGPUComputePipeline old_pipeline = d.pipeline;
+  WGPUBindGroup old_bind_group = d.bind_group;
+  d.pipeline = replacement.pipeline;
+  d.bind_group = replacement.bind_group;
+  if (old_pipeline != nullptr) {
+    wgpuComputePipelineRelease(old_pipeline);
+  }
+  if (old_bind_group != nullptr) {
+    wgpuBindGroupRelease(old_bind_group);
+  }
+}
+
+} // namespace
+
+void invalidate_record(WebGPUGraph& graph) {
+  graph.clear_rms_fusion_site();
+}
+
+void record_rms_norm(
+    WebGPUGraph& graph,
+    int in_id,
+    int weight_id,
+    int out_id,
+    uint32_t num_rows,
+    uint32_t row_width,
+    size_t dispatch_idx,
+    WGPUBuffer params_buf,
+    WGPUBindGroup /*bind_group*/) {
+  Record r = {};
+  r.in_id = in_id;
+  r.weight_id = weight_id;
+  r.out_id = out_id;
+  r.num_rows = num_rows;
+  r.row_width = row_width;
+  r.dispatch_index = dispatch_idx;
+  r.params_buffer = params_buf;
+  graph.offer_rms_fusion_site(std::move(r));
+}
+
+bool try_fuse_add(
+    WebGPUGraph& graph,
+    int in1_id,
+    int in2_id,
+    float alpha,
+    int out_id) {
+  Record r = graph.rms_fusion_site();
+  if (!still_last(graph, r) || r.add_fused) {
+    return false;
+  }
+
+  const bool first_is_rms = (in1_id == r.out_id);
+  const bool second_is_rms = (in2_id == r.out_id);
+  const int resid_id = first_is_rms ? in2_id : in1_id;
+
+  const WebGPUTensor& t_in = graph.get_tensor(r.in_id);
+  const WebGPUTensor& t_w = graph.get_tensor(r.weight_id);
+  const WebGPUTensor& t_out = graph.get_tensor(r.out_id);
+  const WebGPUTensor& t_resid = graph.get_tensor(resid_id);
+  const WebGPUTensor& t_addout = graph.get_tensor(out_id);
+
+  RmsAddCheck c;
+  c.adjacent = true;
+  c.exactly_one_operand_is_rms_out = (first_is_rms != second_is_rms);
+  c.row_width = r.row_width;
+  c.rms_in_numel = numel_of_tensor(t_in);
+  c.rms_out_numel = numel_of_tensor(t_out);
+  c.resid_numel = numel_of_tensor(t_resid);
+  c.add_out_numel = numel_of_tensor(t_addout);
+  c.exact_shape_match = t_in.dims == t_out.dims && t_in.dims == t_resid.dims &&
+      t_in.dims == t_addout.dims;
+  c.all_tensors_fp32 = utils::is_fp32_tensor(t_in) &&
+      utils::is_fp32_tensor(t_w) && utils::is_fp32_tensor(t_out) &&
+      utils::is_fp32_tensor(t_resid) && utils::is_fp32_tensor(t_addout);
+  c.alpha = alpha;
+  c.addout_is_out_buffer = (t_addout.buffer == t_out.buffer);
+  c.addout_is_in_buffer = (t_addout.buffer == t_in.buffer);
+  c.addout_is_weight_buffer = (t_addout.buffer == t_w.buffer);
+  c.resid_is_out_buffer = (t_resid.buffer == t_out.buffer);
+  if (!rms_add_fusable(c)) {
+    return false;
+  }
+  install(
+      graph,
+      r.dispatch_index,
+      kRmsNormVec4AddWGSL,
+      {
+          {0, WGPUBufferBindingType_Storage, t_out.buffer, t_out.nbytes},
+          {1, WGPUBufferBindingType_ReadOnlyStorage, t_in.buffer, t_in.nbytes},
+          {2, WGPUBufferBindingType_ReadOnlyStorage, t_w.buffer, t_w.nbytes},
+          {3, WGPUBufferBindingType_Uniform, r.params_buffer, 16u},
+          {4,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           t_resid.buffer,
+           t_resid.nbytes},
+          {5, WGPUBufferBindingType_Storage, t_addout.buffer, t_addout.nbytes},
+      });
+
+  // The rms_norm resize hook already rewrites params_buf (unchanged 16-byte
+  // layout) and workgroup_count_x, and sets cur_dims(rms out). The add's own
+  // hooks are never registered, so mirror its output extent here. Registered
+  // after the rms hook, so it runs after cur_dims(rms out) has converged.
+  const int rms_in_id = r.in_id;
+  const int fused_resid_id = resid_id;
+  const int addout_id = out_id;
+  auto resize_add = [rms_in_id, fused_resid_id, addout_id](WebGPUGraph& g) {
+    resize_fused_add_outputs(g, rms_in_id, fused_resid_id, addout_id);
+  };
+  graph.add_tensor_resize_hook(rms_in_id, resize_add);
+  if (fused_resid_id != rms_in_id) {
+    graph.add_tensor_resize_hook(fused_resid_id, resize_add);
+  }
+
+  r.add_fused = true;
+  r.resid_id = resid_id;
+  r.addout_id = out_id;
+  graph.offer_rms_fusion_site(std::move(r));
+  return true;
+}
+
+bool try_fuse_scale(WebGPUGraph& graph, int in1_id, int in2_id, int out_id) {
+  Record r = graph.rms_fusion_site();
+  if (!still_last(graph, r) || !r.add_fused) {
+    return false;
+  }
+
+  const bool first_is_add = (in1_id == r.addout_id);
+  const bool second_is_add = (in2_id == r.addout_id);
+  const int scale_id = first_is_add ? in2_id : in1_id;
+
+  const WebGPUTensor& t_in = graph.get_tensor(r.in_id);
+  const WebGPUTensor& t_w = graph.get_tensor(r.weight_id);
+  const WebGPUTensor& t_out = graph.get_tensor(r.out_id);
+  const WebGPUTensor& t_resid = graph.get_tensor(r.resid_id);
+  const WebGPUTensor& t_addout = graph.get_tensor(r.addout_id);
+  const WebGPUTensor& t_scale = graph.get_tensor(scale_id);
+  const WebGPUTensor& t_scaleout = graph.get_tensor(out_id);
+
+  RmsScaleCheck c;
+  c.adjacent_fused_add = true;
+  c.exactly_one_operand_is_add_out = (first_is_add != second_is_add);
+  c.scale_numel = numel_of_tensor(t_scale);
+  c.add_out_numel = numel_of_tensor(t_addout);
+  c.mul_out_numel = numel_of_tensor(t_scaleout);
+  c.all_tensors_fp32 = utils::is_fp32_tensor(t_in) &&
+      utils::is_fp32_tensor(t_w) && utils::is_fp32_tensor(t_out) &&
+      utils::is_fp32_tensor(t_resid) && utils::is_fp32_tensor(t_addout) &&
+      utils::is_fp32_tensor(t_scale) && utils::is_fp32_tensor(t_scaleout);
+  c.scaleout_is_out_buffer = (t_scaleout.buffer == t_out.buffer);
+  c.scaleout_is_in_buffer = (t_scaleout.buffer == t_in.buffer);
+  c.scaleout_is_weight_buffer = (t_scaleout.buffer == t_w.buffer);
+  c.scaleout_is_addout_buffer = (t_scaleout.buffer == t_addout.buffer);
+  c.scaleout_is_resid_buffer = (t_scaleout.buffer == t_resid.buffer);
+  c.scale_is_out_buffer = (t_scale.buffer == t_out.buffer);
+  c.scale_is_addout_buffer = (t_scale.buffer == t_addout.buffer);
+  if (!rms_scale_fusable(c)) {
+    return false;
+  }
+  install(
+      graph,
+      r.dispatch_index,
+      kRmsNormVec4AddScaleWGSL,
+      {
+          {0, WGPUBufferBindingType_Storage, t_out.buffer, t_out.nbytes},
+          {1, WGPUBufferBindingType_ReadOnlyStorage, t_in.buffer, t_in.nbytes},
+          {2, WGPUBufferBindingType_ReadOnlyStorage, t_w.buffer, t_w.nbytes},
+          {3, WGPUBufferBindingType_Uniform, r.params_buffer, 16u},
+          {4,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           t_resid.buffer,
+           t_resid.nbytes},
+          {5, WGPUBufferBindingType_Storage, t_addout.buffer, t_addout.nbytes},
+          {6,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           t_scale.buffer,
+           t_scale.nbytes},
+          {7,
+           WGPUBufferBindingType_Storage,
+           t_scaleout.buffer,
+           t_scaleout.nbytes},
+      });
+
+  const int rms_in_id = r.in_id;
+  const int resid_id = r.resid_id;
+  const int addout_id = r.addout_id;
+  const int fused_scale_id = scale_id;
+  const int scaleout_id = out_id;
+  auto resize_scale =
+      [rms_in_id, resid_id, addout_id, fused_scale_id, scaleout_id](
+          WebGPUGraph& g) {
+        resize_fused_add_outputs(
+            g, rms_in_id, resid_id, addout_id, fused_scale_id, scaleout_id);
+      };
+  graph.add_tensor_resize_hook(rms_in_id, resize_scale);
+  if (resid_id != rms_in_id) {
+    graph.add_tensor_resize_hook(resid_id, resize_scale);
+  }
+  if (fused_scale_id != rms_in_id && fused_scale_id != resid_id) {
+    graph.add_tensor_resize_hook(fused_scale_id, resize_scale);
+  }
+
+  // Nothing further can chain onto this dispatch.
+  graph.clear_rms_fusion_site();
+  return true;
+}
+
+} // namespace fusion
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.h b/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.h
new file mode 100644
index 00000000000..257b120559f
--- /dev/null
+++ b/backends/webgpu/runtime/ops/rms_norm/RmsNormFusion.h
@@ -0,0 +1,210 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+#include <webgpu/webgpu.h>
+
+namespace executorch::backends::webgpu {
+
+class WebGPUGraph;
+
+namespace fusion {
+
+// Dispatch-count reduction for the sealed Gemma4 decode round.
+//
+// et_vk.rms_norm.default is followed, 121 times out of 284, by an
+// aten.add.Tensor that consumes its output with no dispatching op in between
+// and an exactly matching shape; 43 of those adds are then followed by an
+// aten.mul.Tensor against a single-element tensor. Each of those is a separate
+// dispatchWorkgroups over the same 3 x 256 elements the rms_norm workgroup has
+// already loaded.
+//
+// This folds them into the rms_norm dispatch by REWRITING the pipeline and
+// bind group of the dispatch rms_norm_impl just emitted -- it never reorders,
+// never drops a store, and never aliases a buffer, so it is a pure dispatch
+// merge. Every intermediate (t_out, t_addout) is still written, so no consumer
+// of them is assumed dead.
+//
+// The state machine is: record_rms_norm -> try_fuse_add -> try_fuse_scale.
+// Any op that emits a dispatch in between invalidates the record, because
+// every step re-checks that the recorded dispatch is still the LAST one.
+
+// ---------------------------------------------------------------------------
+// Pure fusibility predicates (no GPU, no graph): unit-tested by the guard test.
+// ---------------------------------------------------------------------------
+
+struct RmsAddCheck {
+  // The recorded rms_norm dispatch is still graph.num_dispatches() - 1.
+  bool adjacent = false;
+  // Exactly one of the add's two tensor operands is the rms_norm output.
+  bool exactly_one_operand_is_rms_out = false;
+  uint32_t row_width = 0;
+  uint64_t rms_in_numel = 0;
+  uint64_t rms_out_numel = 0;
+  uint64_t resid_numel = 0;
+  uint64_t add_out_numel = 0;
+  bool exact_shape_match = false;
+  bool all_tensors_fp32 = false;
+  float alpha = 0.0f;
+  // Buffer-identity facts. The fused kernel writes t_out and t_addout and
+  // reads t_in, t_weight and t_resid inside one dispatch, so any overlap that
+  // the two-dispatch form made safe by the inter-dispatch barrier must reject.
+  bool addout_is_out_buffer = false;
+  bool addout_is_in_buffer = false;
+  bool addout_is_weight_buffer = false;
+  bool resid_is_out_buffer = false;
+};
+
+struct RmsScaleCheck {
+  // A rms+add fusion was installed and is still the LAST dispatch.
+  bool adjacent_fused_add = false;
+  // Exactly one of the mul's two operands is the fused add's output.
+  bool exactly_one_operand_is_add_out = false;
+  uint64_t scale_numel = 0; // must be 1 (binary_mul collapses to input2[0])
+  uint64_t add_out_numel = 0;
+  uint64_t mul_out_numel = 0;
+  bool all_tensors_fp32 = false;
+  bool scaleout_is_out_buffer = false;
+  bool scaleout_is_in_buffer = false;
+  bool scaleout_is_weight_buffer = false;
+  bool scaleout_is_addout_buffer = false;
+  bool scaleout_is_resid_buffer = false;
+  bool scale_is_out_buffer = false;
+  bool scale_is_addout_buffer = false;
+};
+
+struct RmsResizeCheck {
+  bool residual_shape_matches = false;
+  bool has_scale = false;
+  uint64_t scale_numel = 0;
+};
+
+// A rejected reason string (nullptr when fusible) makes the guard test able to
+// assert WHICH guard fired, not just that one did. Header-inline so the guard
+// test links without any GPU object.
+inline const char* rms_add_reject_reason(const RmsAddCheck& c) {
+  if (!c.adjacent) {
+    return "rms_norm dispatch is no longer the last one";
+  }
+  if (!c.exactly_one_operand_is_rms_out) {
+    return "add does not consume the rms_norm output exactly once";
+  }
+  if (c.row_width == 0u || c.row_width % 4u != 0u) {
+    return "row_width is not a positive multiple of 4 (vec4 route only)";
+  }
+  if (c.rms_in_numel == 0u || c.rms_in_numel % c.row_width != 0u) {
+    return "rms input numel is not a positive multiple of row_width";
+  }
+  if (c.rms_out_numel != c.rms_in_numel || c.resid_numel != c.rms_in_numel ||
+      c.add_out_numel != c.rms_in_numel) {
+    return "operand numels differ (the fused add is elementwise, not broadcast)";
+  }
+  if (!c.exact_shape_match) {
+    return "operand shapes differ (the fused add does not implement broadcast)";
+  }
+  if (!c.all_tensors_fp32) {
+    return "rms/add tensors are not all fp32";
+  }
+  if (c.alpha != 1.0f) {
+    return "alpha != 1 (the fused kernel folds the add as a bare r + n)";
+  }
+  if (c.addout_is_out_buffer) {
+    return "add output aliases the rms_norm output buffer";
+  }
+  if (c.addout_is_in_buffer) {
+    return "add output aliases the rms_norm input buffer";
+  }
+  if (c.addout_is_weight_buffer) {
+    return "add output aliases the rms_norm weight buffer";
+  }
+  if (c.resid_is_out_buffer) {
+    return "add residual aliases the rms_norm output buffer";
+  }
+  return nullptr;
+}
+
+inline const char* rms_scale_reject_reason(const RmsScaleCheck& c) {
+  if (!c.adjacent_fused_add) {
+    return "no fused rms+add dispatch immediately precedes this mul";
+  }
+  if (!c.exactly_one_operand_is_add_out) {
+    return "mul does not consume the fused add output exactly once";
+  }
+  if (c.scale_numel != 1u) {
+    return "scale operand is not a single element";
+  }
+  if (c.mul_out_numel != c.add_out_numel) {
+    return "mul output numel differs from the add output numel";
+  }
+  if (!c.all_tensors_fp32) {
+    return "rms/add/scale tensors are not all fp32";
+  }
+  if (c.scaleout_is_out_buffer || c.scaleout_is_in_buffer ||
+      c.scaleout_is_weight_buffer || c.scaleout_is_addout_buffer ||
+      c.scaleout_is_resid_buffer) {
+    return "mul output aliases a buffer the fused kernel already touches";
+  }
+  if (c.scale_is_out_buffer || c.scale_is_addout_buffer) {
+    return "scale operand aliases a buffer the fused kernel writes";
+  }
+  return nullptr;
+}
+
+inline const char* rms_resize_reject_reason(const RmsResizeCheck& c) {
+  if (!c.residual_shape_matches) {
+    return "live residual shape requires broadcast";
+  }
+  if (c.has_scale && c.scale_numel != 1u) {
+    return "live scale is not scalar";
+  }
+  return nullptr;
+}
+
+inline bool rms_add_fusable(const RmsAddCheck& c) {
+  return rms_add_reject_reason(c) == nullptr;
+}
+inline bool rms_scale_fusable(const RmsScaleCheck& c) {
+  return rms_scale_reject_reason(c) == nullptr;
+}
+
+// ---------------------------------------------------------------------------
+// Build-time hooks called from the three op handlers.
+// ---------------------------------------------------------------------------
+
+// Called by rms_norm_impl right after it emits its dispatch, only for the vec4
+// route (row_width % 4 == 0) at the fixed 64-wide workgroup.
+void record_rms_norm(
+    WebGPUGraph& graph,
+    int in_id,
+    int weight_id,
+    int out_id,
+    uint32_t num_rows,
+    uint32_t row_width,
+    size_t dispatch_idx,
+    WGPUBuffer params_buf,
+    WGPUBindGroup bind_group);
+
+// Invalidate the record (any handler that emits a dispatch of its own).
+void invalidate_record(WebGPUGraph& graph);
+
+// Called at the top of add_impl / mul_impl. Returns true when the op has been
+// folded into the preceding dispatch and the caller must emit nothing.
+bool try_fuse_add(
+    WebGPUGraph& graph,
+    int in1_id,
+    int in2_id,
+    float alpha,
+    int out_id);
+bool try_fuse_scale(WebGPUGraph& graph, int in1_id, int in2_id, int out_id);
+
+} // namespace fusion
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl
index 11302434ec2..101f2f5b0a0 100644
--- a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl
@@ -39,8 +39,10 @@ $if VEC == 4:
 @compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>) {
-  let row_idx = wid.x;
+  // 2D-fold: rows can exceed the 65535 per-dim cap (QK-norm at prefill).
+  let row_idx = wid.x + wid.y * num_workgroups.x;
   let worker_id = lid.x;
 
   if (row_idx >= params.num_rows) {
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add.wgsl
new file mode 100644
index 00000000000..eef0a2df81a
--- /dev/null
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add.wgsl
@@ -0,0 +1,107 @@
+// 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<vec4<f32>>;
+@group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> t_weight: array<vec4<f32>>;
+
+// Byte-identical to the unfused rms_norm_vec4 Params, so RmsNorm.cpp's existing
+// resize hook keeps rewriting the SAME uniform buffer with the SAME 16 bytes.
+struct Params {
+  num_rows: u32,
+  row_width: u32,
+  epsilon: f32,
+  _pad: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+@group(0) @binding(4) var<storage, read> t_resid: array<vec4<f32>>;
+@group(0) @binding(5) var<storage, read_write> t_addout: array<vec4<f32>>;
+
+// Fixed workgroup size (Apple scalar ALU: no override loop bounds). 64 is the
+// value clamp_workgroup_size_pow2 already yields for the unfused kernel, so the
+// tree reduction visits the identical strides in the identical order and the
+// normalized result is bit-identical.
+const wg_size: u32 = 64u;
+
+var<workgroup> shared_sum: array<f32, 64>;
+
+fn reduce_shared(worker_id: u32) {
+  workgroupBarrier();
+  var stride: u32 = wg_size / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (worker_id < stride) {
+      shared_sum[worker_id] = shared_sum[worker_id] + shared_sum[worker_id + stride];
+    }
+    workgroupBarrier();
+    stride = stride >> 1u;
+  }
+}
+
+// rms_norm_vec4 with the immediately following aten.add.Tensor folded in.
+// The reduction and the normalized store are byte-for-byte the unfused vec4
+// kernel; the add evaluates the same expression binary_add.wgsl does, at the
+// same element index. The host only takes this route when alpha == 1.0, which
+// makes `fuse.alpha * n` exact under any fma contraction the backend chooses.
+// t_out is still written, so no consumer of it is assumed dead.
+// The host guards alpha == 1, and IEEE addition is commutative, so both
+// binary_add operand orders (resid + 1.0*n and n + 1.0*resid) collapse to the
+// single expression `r + n` -- bit-identically, and under any fma contraction,
+// because fma(1, n, r) == n + r exactly. There is therefore no alpha and no
+// operand-order flag to carry.
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-fold: rows can exceed the 65535 per-dim cap (QK-norm at prefill).
+  let row_idx = wid.x + wid.y * num_workgroups.x;
+  let worker_id = lid.x;
+
+  if (row_idx >= params.num_rows) {
+    return;
+  }
+
+  let rw4 = params.row_width / 4u;
+  let base4 = row_idx * rw4;
+
+  var local_sq_sum: f32 = 0.0;
+  var x4: u32 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let v = t_in[base4 + x4];
+    local_sq_sum = local_sq_sum + dot(v, v);
+    x4 = x4 + wg_size;
+  }
+
+  shared_sum[worker_id] = local_sq_sum;
+  reduce_shared(worker_id);
+
+  let mean_sq = shared_sum[0] / f32(params.row_width);
+  let rstd = inverseSqrt(mean_sq + params.epsilon);
+
+  x4 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let r = t_resid[base4 + x4];
+    let n = t_in[base4 + x4] * rstd * t_weight[x4];
+    t_out[base4 + x4] = n;
+    t_addout[base4 + x4] = r + n;
+    x4 = x4 + wg_size;
+  }
+}
+// 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.
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale.wgsl
new file mode 100644
index 00000000000..189ac842779
--- /dev/null
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale.wgsl
@@ -0,0 +1,113 @@
+// 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<vec4<f32>>;
+@group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> t_weight: array<vec4<f32>>;
+
+// Byte-identical to the unfused rms_norm_vec4 Params, so RmsNorm.cpp's existing
+// resize hook keeps rewriting the SAME uniform buffer with the SAME 16 bytes.
+struct Params {
+  num_rows: u32,
+  row_width: u32,
+  epsilon: f32,
+  _pad: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+@group(0) @binding(4) var<storage, read> t_resid: array<vec4<f32>>;
+@group(0) @binding(5) var<storage, read_write> t_addout: array<vec4<f32>>;
+
+@group(0) @binding(6) var<storage, read> t_scale: array<f32>;
+@group(0) @binding(7) var<storage, read_write> t_scaleout: array<vec4<f32>>;
+
+// Fixed workgroup size (Apple scalar ALU: no override loop bounds). 64 is the
+// value clamp_workgroup_size_pow2 already yields for the unfused kernel, so the
+// tree reduction visits the identical strides in the identical order and the
+// normalized result is bit-identical.
+const wg_size: u32 = 64u;
+
+var<workgroup> shared_sum: array<f32, 64>;
+
+fn reduce_shared(worker_id: u32) {
+  workgroupBarrier();
+  var stride: u32 = wg_size / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (worker_id < stride) {
+      shared_sum[worker_id] = shared_sum[worker_id] + shared_sum[worker_id + stride];
+    }
+    workgroupBarrier();
+    stride = stride >> 1u;
+  }
+}
+
+// rms_norm_vec4 + the following aten.add.Tensor + the following
+// aten.mul.Tensor by a single-element tensor (binary_mul's broadcast path
+// collapses that operand to input2[0]). All three stores are kept, so no
+// consumer liveness assumption is made: this is a dispatch merge, not a graph
+// rewrite.
+// The host guards alpha == 1, and IEEE addition is commutative, so both
+// binary_add operand orders (resid + 1.0*n and n + 1.0*resid) collapse to the
+// single expression `r + n` -- bit-identically, and under any fma contraction,
+// because fma(1, n, r) == n + r exactly. There is therefore no alpha and no
+// operand-order flag to carry.
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-fold: rows can exceed the 65535 per-dim cap (QK-norm at prefill).
+  let row_idx = wid.x + wid.y * num_workgroups.x;
+  let worker_id = lid.x;
+
+  if (row_idx >= params.num_rows) {
+    return;
+  }
+
+  let rw4 = params.row_width / 4u;
+  let base4 = row_idx * rw4;
+
+  var local_sq_sum: f32 = 0.0;
+  var x4: u32 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let v = t_in[base4 + x4];
+    local_sq_sum = local_sq_sum + dot(v, v);
+    x4 = x4 + wg_size;
+  }
+
+  shared_sum[worker_id] = local_sq_sum;
+  reduce_shared(worker_id);
+
+  let mean_sq = shared_sum[0] / f32(params.row_width);
+  let rstd = inverseSqrt(mean_sq + params.epsilon);
+
+  let s = t_scale[0];
+
+  x4 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let r = t_resid[base4 + x4];
+    let n = t_in[base4 + x4] * rstd * t_weight[x4];
+    t_out[base4 + x4] = n;
+    let a = r + n;
+    t_addout[base4 + x4] = a;
+    t_scaleout[base4 + x4] = a * s;
+    x4 = x4 + wg_size;
+  }
+}
+// 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.
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale_wgsl.h
new file mode 100644
index 00000000000..029b480a9e2
--- /dev/null
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_scale_wgsl.h
@@ -0,0 +1,137 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from rms_norm_vec4_add_scale.wgsl - DO NOT EDIT.
+// wgsl-sha256: 13387491431ae152eedd2f3e7e18cf4b7cd1bc24d6c6df4607efe8f2eb97eb1a
+inline constexpr const char* kRmsNormVec4AddScaleWGSL = R"(
+// 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<vec4<f32>>;
+@group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> t_weight: array<vec4<f32>>;
+
+// Byte-identical to the unfused rms_norm_vec4 Params, so RmsNorm.cpp's existing
+// resize hook keeps rewriting the SAME uniform buffer with the SAME 16 bytes.
+struct Params {
+  num_rows: u32,
+  row_width: u32,
+  epsilon: f32,
+  _pad: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+@group(0) @binding(4) var<storage, read> t_resid: array<vec4<f32>>;
+@group(0) @binding(5) var<storage, read_write> t_addout: array<vec4<f32>>;
+
+@group(0) @binding(6) var<storage, read> t_scale: array<f32>;
+@group(0) @binding(7) var<storage, read_write> t_scaleout: array<vec4<f32>>;
+
+// Fixed workgroup size (Apple scalar ALU: no override loop bounds). 64 is the
+// value clamp_workgroup_size_pow2 already yields for the unfused kernel, so the
+// tree reduction visits the identical strides in the identical order and the
+// normalized result is bit-identical.
+const wg_size: u32 = 64u;
+
+var<workgroup> shared_sum: array<f32, 64>;
+
+fn reduce_shared(worker_id: u32) {
+  workgroupBarrier();
+  var stride: u32 = wg_size / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (worker_id < stride) {
+      shared_sum[worker_id] = shared_sum[worker_id] + shared_sum[worker_id + stride];
+    }
+    workgroupBarrier();
+    stride = stride >> 1u;
+  }
+}
+
+// rms_norm_vec4 + the following aten.add.Tensor + the following
+// aten.mul.Tensor by a single-element tensor (binary_mul's broadcast path
+// collapses that operand to input2[0]). All three stores are kept, so no
+// consumer liveness assumption is made: this is a dispatch merge, not a graph
+// rewrite.
+// The host guards alpha == 1, and IEEE addition is commutative, so both
+// binary_add operand orders (resid + 1.0*n and n + 1.0*resid) collapse to the
+// single expression `r + n` -- bit-identically, and under any fma contraction,
+// because fma(1, n, r) == n + r exactly. There is therefore no alpha and no
+// operand-order flag to carry.
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-fold: rows can exceed the 65535 per-dim cap (QK-norm at prefill).
+  let row_idx = wid.x + wid.y * num_workgroups.x;
+  let worker_id = lid.x;
+
+  if (row_idx >= params.num_rows) {
+    return;
+  }
+
+  let rw4 = params.row_width / 4u;
+  let base4 = row_idx * rw4;
+
+  var local_sq_sum: f32 = 0.0;
+  var x4: u32 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let v = t_in[base4 + x4];
+    local_sq_sum = local_sq_sum + dot(v, v);
+    x4 = x4 + wg_size;
+  }
+
+  shared_sum[worker_id] = local_sq_sum;
+  reduce_shared(worker_id);
+
+  let mean_sq = shared_sum[0] / f32(params.row_width);
+  let rstd = inverseSqrt(mean_sq + params.epsilon);
+
+  let s = t_scale[0];
+
+  x4 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let r = t_resid[base4 + x4];
+    let n = t_in[base4 + x4] * rstd * t_weight[x4];
+    t_out[base4 + x4] = n;
+    let a = r + n;
+    t_addout[base4 + x4] = a;
+    t_scaleout[base4 + x4] = a * s;
+    x4 = x4 + wg_size;
+  }
+}
+// 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.
+)";
+
+inline constexpr uint32_t kRmsNormVec4AddScaleWorkgroupSizeX = 64;
+inline constexpr uint32_t kRmsNormVec4AddScaleWorkgroupSizeY = 1;
+inline constexpr uint32_t kRmsNormVec4AddScaleWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_wgsl.h
new file mode 100644
index 00000000000..d4126351bbc
--- /dev/null
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_add_wgsl.h
@@ -0,0 +1,131 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from rms_norm_vec4_add.wgsl - DO NOT EDIT.
+// wgsl-sha256: d4855cb6c90093ba216ae4a877567a3f70a5ec0d4b439f4fc012bedf24fc2027
+inline constexpr const char* kRmsNormVec4AddWGSL = R"(
+// 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<vec4<f32>>;
+@group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> t_weight: array<vec4<f32>>;
+
+// Byte-identical to the unfused rms_norm_vec4 Params, so RmsNorm.cpp's existing
+// resize hook keeps rewriting the SAME uniform buffer with the SAME 16 bytes.
+struct Params {
+  num_rows: u32,
+  row_width: u32,
+  epsilon: f32,
+  _pad: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+@group(0) @binding(4) var<storage, read> t_resid: array<vec4<f32>>;
+@group(0) @binding(5) var<storage, read_write> t_addout: array<vec4<f32>>;
+
+// Fixed workgroup size (Apple scalar ALU: no override loop bounds). 64 is the
+// value clamp_workgroup_size_pow2 already yields for the unfused kernel, so the
+// tree reduction visits the identical strides in the identical order and the
+// normalized result is bit-identical.
+const wg_size: u32 = 64u;
+
+var<workgroup> shared_sum: array<f32, 64>;
+
+fn reduce_shared(worker_id: u32) {
+  workgroupBarrier();
+  var stride: u32 = wg_size / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (worker_id < stride) {
+      shared_sum[worker_id] = shared_sum[worker_id] + shared_sum[worker_id + stride];
+    }
+    workgroupBarrier();
+    stride = stride >> 1u;
+  }
+}
+
+// rms_norm_vec4 with the immediately following aten.add.Tensor folded in.
+// The reduction and the normalized store are byte-for-byte the unfused vec4
+// kernel; the add evaluates the same expression binary_add.wgsl does, at the
+// same element index. The host only takes this route when alpha == 1.0, which
+// makes `fuse.alpha * n` exact under any fma contraction the backend chooses.
+// t_out is still written, so no consumer of it is assumed dead.
+// The host guards alpha == 1, and IEEE addition is commutative, so both
+// binary_add operand orders (resid + 1.0*n and n + 1.0*resid) collapse to the
+// single expression `r + n` -- bit-identically, and under any fma contraction,
+// because fma(1, n, r) == n + r exactly. There is therefore no alpha and no
+// operand-order flag to carry.
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-fold: rows can exceed the 65535 per-dim cap (QK-norm at prefill).
+  let row_idx = wid.x + wid.y * num_workgroups.x;
+  let worker_id = lid.x;
+
+  if (row_idx >= params.num_rows) {
+    return;
+  }
+
+  let rw4 = params.row_width / 4u;
+  let base4 = row_idx * rw4;
+
+  var local_sq_sum: f32 = 0.0;
+  var x4: u32 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let v = t_in[base4 + x4];
+    local_sq_sum = local_sq_sum + dot(v, v);
+    x4 = x4 + wg_size;
+  }
+
+  shared_sum[worker_id] = local_sq_sum;
+  reduce_shared(worker_id);
+
+  let mean_sq = shared_sum[0] / f32(params.row_width);
+  let rstd = inverseSqrt(mean_sq + params.epsilon);
+
+  x4 = worker_id;
+  loop {
+    if (x4 >= rw4) {
+      break;
+    }
+    let r = t_resid[base4 + x4];
+    let n = t_in[base4 + x4] * rstd * t_weight[x4];
+    t_out[base4 + x4] = n;
+    t_addout[base4 + x4] = r + n;
+    x4 = x4 + wg_size;
+  }
+}
+// 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.
+)";
+
+inline constexpr uint32_t kRmsNormVec4AddWorkgroupSizeX = 64;
+inline constexpr uint32_t kRmsNormVec4AddWorkgroupSizeY = 1;
+inline constexpr uint32_t kRmsNormVec4AddWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h
index 2e17ac02aa0..5acec33e8ff 100644
--- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from rms_norm.wgsl - DO NOT EDIT.
-// wgsl-sha256: 62fdfe03fc67eb44fa17ca5e91e433b3a45a96c76151777856d6f557fd829919
+// wgsl-sha256: 7c8ddd078f5c53d9ff6823b11244fd0f23b93373e2eff80601925e7ac9c753ed
 inline constexpr const char* kRmsNormVec4WGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<vec4<f32>>;
 @group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
@@ -53,8 +53,10 @@ fn reduce_shared(worker_id: u32) {
 @compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>) {
-  let row_idx = wid.x;
+  // 2D-fold: rows can exceed the 65535 per-dim cap (QK-norm at prefill).
+  let row_idx = wid.x + wid.y * num_workgroups.x;
   let worker_id = lid.x;
 
   if (row_idx >= params.num_rows) {
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h
index 841b146b2ef..a8cb6206967 100644
--- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from rms_norm.wgsl - DO NOT EDIT.
-// wgsl-sha256: 55bc862d64451f5fe8f530114f0318dfbc724017314091418a1fe41b180be7c6
+// wgsl-sha256: 94e087f05ff901b76bc272959771a23e3452e489da8f2704f5e78df6d730c2ca
 inline constexpr const char* kRmsNormWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_in: array<f32>;
@@ -50,8 +50,10 @@ fn reduce_shared(worker_id: u32) {
 @compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>) {
-  let row_idx = wid.x;
+  // 2D-fold: rows can exceed the 65535 per-dim cap (QK-norm at prefill).
+  let row_idx = wid.x + wid.y * num_workgroups.x;
   let worker_id = lid.x;
 
   if (row_idx >= params.num_rows) {
diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp
index 29f00c5823f..5e19405d6f3 100644
--- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp
+++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp
@@ -15,6 +15,7 @@
 
 #include <cstdint>
 #include <stdexcept>
+#include <string>
 
 namespace executorch::backends::webgpu {
 
@@ -630,12 +631,197 @@ void apply_rotary_emb_hf_impl(
   }
 }
 
+struct RotaryHfSingleResizeContext {
+  int x_id;
+  int out_id;
+  int start_pos_id;
+  bool dynamic_pos;
+  uint32_t baked_start_pos;
+  uint32_t batch;
+  uint32_t n_heads;
+  uint32_t head_dim;
+  uint32_t max_seq;
+  WGPUBuffer uniform;
+};
+
+void resize_rope_hf_single(
+    WebGPUGraph& graph,
+    const RotaryHfSingleResizeContext& context) {
+  const auto& dims = graph.cur_dims(context.x_id);
+  if (dims.size() != 4 || dims[0] != context.batch || dims[1] <= 0 ||
+      static_cast<uint64_t>(dims[1]) > UINT32_MAX ||
+      dims[2] != context.n_heads || dims[3] != context.head_dim) {
+    throw std::runtime_error(
+        "apply_rotary_emb_hf_single(resize): invalid BSHD shape");
+  }
+  const uint32_t seq = static_cast<uint32_t>(dims[1]);
+  const uint64_t numel = utils::numel_of(dims);
+  if (numel == 0 || numel / 2u > UINT32_MAX) {
+    throw std::runtime_error(
+        "apply_rotary_emb_hf_single(resize): pair count exceeds uint32");
+  }
+
+  uint32_t start_pos = context.baked_start_pos;
+  if (context.dynamic_pos) {
+    const int64_t pos = graph.read_symint(context.start_pos_id);
+    if (pos < 0 || static_cast<uint64_t>(pos) > UINT32_MAX) {
+      throw std::runtime_error(
+          "apply_rotary_emb_hf_single(resize): invalid start_pos");
+    }
+    start_pos = static_cast<uint32_t>(pos);
+  }
+  if (static_cast<uint64_t>(start_pos) + seq > context.max_seq) {
+    throw std::runtime_error(
+        "apply_rotary_emb_hf_single(resize): position range exceeds freqs");
+  }
+
+  const RotaryHfParams params = {
+      context.n_heads,
+      seq,
+      context.head_dim,
+      context.head_dim / 2u,
+      static_cast<uint32_t>(numel / 2u),
+      context.head_dim,
+      start_pos,
+      0u};
+  wgpuQueueWriteBuffer(
+      graph.queue(), context.uniform, 0, &params, sizeof(params));
+  graph.set_cur_dims(context.out_id, dims);
+}
+
+void apply_rotary_emb_hf_single_impl(
+    WebGPUGraph& graph,
+    const std::vector<int>& args) {
+  if (args.size() != 5) {
+    throw std::runtime_error(
+        "WebGPU apply_rotary_emb_hf_single: expected 5 args");
+  }
+  const int x_id = args.at(0);
+  const int freqs_cos_id = args.at(1);
+  const int freqs_sin_id = args.at(2);
+  const int start_pos_id = args.at(3);
+  const int out_id = args.at(4);
+  const auto& x = graph.get_tensor(x_id);
+  const auto& freqs_cos = graph.get_tensor(freqs_cos_id);
+  const auto& freqs_sin = graph.get_tensor(freqs_sin_id);
+  const auto& out = graph.get_tensor(out_id);
+
+  if (x.dims.size() != 4 || out.dims != x.dims ||
+      freqs_cos.dims.size() != 2 || freqs_sin.dims != freqs_cos.dims) {
+    throw std::runtime_error(
+        "WebGPU apply_rotary_emb_hf_single: expected x/out [B,S,H,D] and "
+        "matching freqs [max_seq,D]");
+  }
+  const auto positive_u32 = [](int64_t value, const char* label) {
+    if (value <= 0 || static_cast<uint64_t>(value) > UINT32_MAX) {
+      throw std::runtime_error(
+          std::string("WebGPU apply_rotary_emb_hf_single: invalid ") + label);
+    }
+    return static_cast<uint32_t>(value);
+  };
+  const uint32_t batch = positive_u32(x.dims[0], "batch");
+  const uint32_t seq = positive_u32(x.dims[1], "sequence length");
+  const uint32_t n_heads = positive_u32(x.dims[2], "head count");
+  const uint32_t head_dim = positive_u32(x.dims[3], "head dimension");
+  const uint32_t max_seq =
+      positive_u32(freqs_cos.dims[0], "frequency row count");
+  const uint32_t rotary_dim =
+      positive_u32(freqs_cos.dims[1], "rotary dimension");
+  if (head_dim % 2u != 0u || rotary_dim != head_dim) {
+    throw std::runtime_error(
+        "WebGPU apply_rotary_emb_hf_single: requires full even head_dim");
+  }
+
+  const WebGPUTensor* tensors[] = {&x, &freqs_cos, &freqs_sin, &out};
+  for (const WebGPUTensor* tensor : tensors) {
+    if (!utils::is_fp32_tensor(*tensor)) {
+      throw std::runtime_error(
+          "WebGPU apply_rotary_emb_hf_single: all tensors must be fp32");
+    }
+  }
+  const uint64_t x_numel = utils::numel_of(x.dims);
+  const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims);
+  if (x_numel / 2u > UINT32_MAX ||
+      freqs_numel != static_cast<uint64_t>(max_seq) * head_dim) {
+    throw std::runtime_error(
+        "WebGPU apply_rotary_emb_hf_single: invalid tensor extent");
+  }
+
+  const auto start_pos_type = graph.get_value_type(start_pos_id);
+  const bool dynamic_pos =
+      start_pos_type == WebGPUGraph::ValueType::SymInt;
+  int64_t start_pos;
+  if (dynamic_pos) {
+    start_pos = graph.read_symint(start_pos_id);
+  } else if (start_pos_type == WebGPUGraph::ValueType::Int) {
+    start_pos = graph.get_int(start_pos_id);
+  } else {
+    throw std::runtime_error(
+        "WebGPU apply_rotary_emb_hf_single: start_pos must be Int or SymInt");
+  }
+  if (start_pos < 0 || static_cast<uint64_t>(start_pos) > UINT32_MAX ||
+      static_cast<uint64_t>(start_pos) + seq > max_seq) {
+    throw std::runtime_error(
+        "WebGPU apply_rotary_emb_hf_single: position range exceeds freqs");
+  }
+
+  const uint32_t wg_size = utils::clamp_workgroup_size(
+      graph.device(), get_webgpu_shader_info(kRotaryHfShader).workgroup_size_x);
+  const RopeGridContext grid_context = {
+      x_id,
+      wg_size,
+      RopeGridPolicy::FoldedTwoDimensional,
+      "apply_rotary_emb_hf_single"};
+  (void)pick_rope_grid(graph, grid_context);
+  const RotaryHfParams params = {
+      n_heads,
+      seq,
+      head_dim,
+      head_dim / 2u,
+      static_cast<uint32_t>(x_numel / 2u),
+      rotary_dim,
+      static_cast<uint32_t>(start_pos),
+      0u};
+  const WGPUBuffer uniform = add_rope_dispatch(
+      graph,
+      kRotaryHfShader,
+      "apply_rotary_emb_hf_single",
+      x,
+      out,
+      freqs_cos,
+      freqs_sin,
+      params,
+      x_id,
+      grid_context,
+      wg_size);
+
+  const RotaryHfSingleResizeContext resize_context = {
+      x_id,
+      out_id,
+      start_pos_id,
+      dynamic_pos,
+      static_cast<uint32_t>(start_pos),
+      batch,
+      n_heads,
+      head_dim,
+      max_seq,
+      uniform};
+  graph.add_tensor_resize_hook(x_id, resize_rope_hf_single, resize_context);
+  if (dynamic_pos) {
+    graph.add_resize_hook(
+        start_pos_id, resize_rope_hf_single, resize_context);
+  }
+}
+
 } // namespace
 
 WEBGPU_REGISTER_OPERATORS {
   WEBGPU_REGISTER_OP(et_vk.apply_rotary_emb.default, apply_rotary_emb_impl);
   WEBGPU_REGISTER_OP(
       et_vk.apply_rotary_emb_hf.default, apply_rotary_emb_hf_impl);
+  WEBGPU_REGISTER_OP(
+      et_vk.apply_rotary_emb_hf_single.default,
+      apply_rotary_emb_hf_single_impl);
 }
 
 } // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp b/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp
index be333c99d53..3d1be020eba 100644
--- a/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp
+++ b/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp
@@ -9,7 +9,6 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 
-#include <algorithm>
 #include <cstdint>
 #include <functional>
 #include <stdexcept>
@@ -28,10 +27,9 @@ void select_as_symint_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   if (graph.get_value_type(out_id) != WebGPUGraph::ValueType::SymInt) {
     throw std::runtime_error("select_as_symint: output is not a SymInt");
   }
-  const std::vector<int>& inputs = graph.input_ids();
-  if (std::find(inputs.begin(), inputs.end(), x_id) == inputs.end()) {
-    throw std::runtime_error(
-        "select_as_symint: source tensor is not a graph input");
+  if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int ||
+      graph.get_value_type(index_id) != WebGPUGraph::ValueType::Int) {
+    throw std::runtime_error("select_as_symint: dim/index must be static Ints");
   }
   graph.add_symint_source(
       out_id,
diff --git a/backends/webgpu/runtime/ops/slice/Slice.cpp b/backends/webgpu/runtime/ops/slice/Slice.cpp
index f4fcb6088a8..0a3ae33a142 100644
--- a/backends/webgpu/runtime/ops/slice/Slice.cpp
+++ b/backends/webgpu/runtime/ops/slice/Slice.cpp
@@ -10,6 +10,9 @@
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/TensorMeta.h>
+#include <executorch/backends/webgpu/runtime/ops/slice/SliceDispatch.h>
+#include <executorch/backends/webgpu/runtime/ops/slice/slice_dual_guard.h>
+#include <executorch/backends/webgpu/runtime/ops/slice/slice_dual_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/slice/slice_wgsl.h>
 
 #include <webgpu/webgpu.h>
@@ -107,6 +110,156 @@ int64_t norm_clamp(int64_t idx, int64_t size) {
   return idx < 0 ? 0 : (idx > size ? size : idx);
 }
 
+// ---------------------------------------------------------------------------
+// Dual-store dispatch merge.
+//
+// In the sealed Gemma4 decode graph, 67 of the 250 aten.slice_copy.Tensor calls
+// take the output of the IMMEDIATELY preceding slice_copy and re-slice it over
+// its whole extent -- a pure elementwise copy, 126 of the 250 slices are such
+// full-span copies overall. Instead of a second gather dispatch, the preceding
+// slice's dispatch is re-bound to store its gathered value to BOTH
+// destinations. The first destination is still written, so no consumer of it is
+// assumed dead, nothing is reordered, and no buffer is aliased: it is a pure
+// dispatch merge.
+// ---------------------------------------------------------------------------
+
+// True when read_index would return a value rather than throw. Eligibility is
+// a "no" here, never a build-time throw: `end` is otherwise read only by the
+// resize hook, which never runs on a static graph.
+bool is_static_index(WebGPUGraph& graph, int id) {
+  switch (graph.get_value_type(id)) {
+    case WebGPUGraph::ValueType::Int:
+    case WebGPUGraph::ValueType::Null:
+      return true;
+    case WebGPUGraph::ValueType::Double: {
+      const double d = graph.get_double(id);
+      return !std::isnan(d) && d >= -9223372036854775808.0 &&
+          d < 9223372036854775808.0 &&
+          static_cast<double>(static_cast<int64_t>(d)) == d;
+    }
+    default:
+      return false;
+  }
+}
+
+// Resolve the slice bounds at the serialized maximum shape and hand them to the
+// pure slice_dual_full_span predicate (unit-tested by slice_dual_guard_test).
+bool is_static_full_span(
+    WebGPUGraph& graph,
+    int start_id,
+    int end_id,
+    int64_t step,
+    int64_t dim_size) {
+  SliceDualSpan s;
+  s.step = step;
+  s.start_is_symint =
+      graph.get_value_type(start_id) == WebGPUGraph::ValueType::SymInt;
+  s.end_is_symint =
+      graph.get_value_type(end_id) == WebGPUGraph::ValueType::SymInt;
+  s.dim_size = dim_size;
+  if (!s.start_is_symint) {
+    if (!is_static_index(graph, start_id)) {
+      return false;
+    }
+    s.start = read_index(graph, start_id, 0);
+  }
+  if (!s.end_is_symint) {
+    if (!is_static_index(graph, end_id)) {
+      return false;
+    }
+    s.end = read_index(graph, end_id, dim_size);
+  }
+  return slice_dual_full_span(s);
+}
+
+// Re-bind the recorded dispatch to slice_dual.wgsl with `out2` appended.
+// Returns false (leaving the graph untouched) unless every guard holds.
+bool try_dual_store(
+    WebGPUGraph& graph,
+    int in_id,
+    int out_id,
+    int start_id,
+    int end_id,
+    int64_t dim,
+    int64_t step) {
+  const WebGPUGraph::SliceChain& c = graph.slice_chain();
+  if (!c.valid || c.out_id != in_id ||
+      c.dispatch_idx + 1 != graph.num_dispatches()) {
+    return false;
+  }
+
+  const auto& in_tensor = graph.get_tensor(in_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+  // Pure copy: identical extents on every dim, full static span on the sliced
+  // dim, and a distinct destination buffer from both operands of the gather.
+  if (in_tensor.dims != out_tensor.dims ||
+      in_tensor.nbytes != out_tensor.nbytes) {
+    return false;
+  }
+  if (!is_static_full_span(
+          graph, start_id, end_id, step, in_tensor.dims[dim])) {
+    return false;
+  }
+  if (!slice_dual_buffers_ok(out_tensor.buffer, c.out_buffer, c.in_buffer)) {
+    return false;
+  }
+  uint64_t out_numel = 1;
+  for (int64_t d : out_tensor.dims) {
+    out_numel *= static_cast<uint64_t>(d);
+  }
+  if (out_tensor.nbytes != out_numel * sizeof(float)) {
+    return false;
+  }
+
+  WGPUDevice device = graph.device();
+
+  // No override constants: slice_dual declares @workgroup_size(64, 1, 1).
+  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+      device,
+      kSliceDualWGSL,
+      {
+          {0, WGPUBufferBindingType_ReadOnlyStorage, c.in_buffer, c.in_nbytes},
+          {1, WGPUBufferBindingType_Storage, c.out_buffer, c.out_nbytes},
+          {2,
+           WGPUBufferBindingType_Uniform,
+           c.out_meta_buf,
+           sizeof(TensorMeta)},
+          {3, WGPUBufferBindingType_Uniform, c.in_meta_buf, sizeof(TensorMeta)},
+          {4, WGPUBufferBindingType_Uniform, c.params_buf, sizeof(SliceParams)},
+          {5,
+           WGPUBufferBindingType_Storage,
+           out_tensor.buffer,
+           out_tensor.nbytes},
+      });
+
+  // The graph owns a dispatch's pipeline/bind group (released in its dtor), so
+  // the ones being replaced are released here. bundle.pipeline/bind_group are
+  // deliberately NOT released by ~ComputePipelineBundle: they move to the
+  // dispatch, exactly as add_dispatch takes them everywhere else.
+  WebGPUDispatch& d = graph.dispatch_at(c.dispatch_idx);
+  if (d.pipeline) {
+    wgpuComputePipelineRelease(d.pipeline);
+  }
+  if (d.bind_group) {
+    wgpuBindGroupRelease(d.bind_group);
+  }
+  d.pipeline = bundle.pipeline;
+  d.bind_group = bundle.bind_group;
+
+  // The recorded slice's own hook already rewrites the metas/params and the
+  // workgroup count and sets cur_dims(in_id). This mirrors the extent onto the
+  // second destination -- the exact trigger and result the skipped slice's own
+  // recompute would have produced, registered later so it runs after it.
+  graph.add_tensor_resize_hook(in_id, [in_id, out_id](WebGPUGraph& g) {
+    g.set_cur_dims(out_id, g.cur_dims(in_id));
+  });
+
+  // slice_dual has exactly two destinations; a third chained copy must not try
+  // to attach to this dispatch.
+  graph.clear_slice_chain();
+  return true;
+}
+
 void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   // args: [self, dim, start, end, step, out]. start/end may be dynamic SymInts;
   // a resize hook recomputes the live extent on `dim` (out[dim] / cur_dims).
@@ -151,8 +304,15 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   params.start = static_cast<uint32_t>(start);
   params.step = static_cast<uint32_t>(step);
 
+  // Dispatch merge: when this slice is a whole-extent copy of the slice
+  // emitted immediately before it, re-bind that dispatch to store into this
+  // output too instead of emitting a second gather.
+  if (try_dual_store(graph, in_id, out_id, start_id, end_id, dim, step)) {
+    return;
+  }
+
   uint32_t wg_size = utils::clamp_workgroup_size(device, kSliceWorkgroupSizeX);
-  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+  utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
       device, out_meta.numel, wg_size, "slice");
 
   WGPUConstantEntry wg_size_constant = {};
@@ -186,8 +346,9 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       &wg_size_constant,
       1);
 
-  graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
-  const size_t dispatch_idx = graph.num_dispatches() - 1;
+  const size_t dispatch_idx =
+      graph.add_dispatch({bundle.pipeline, bundle.bind_group, 1u, "slice"});
+  set_slice_dispatch_grid(graph, dispatch_idx, workgroup_count);
 
   // Dynamic shapes: live start/end -> out[dim] len + meta/params/dispatch.
   auto recompute = [in_id,
@@ -228,9 +389,9 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     p.start = static_cast<uint32_t>(start);
     p.step = static_cast<uint32_t>(step);
     wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
-    g.dispatch_at(dispatch_idx).workgroup_count_x =
-        utils::compute_1d_workgroup_count(
-            g.device(), om.numel, wg_size, "slice(resize)");
+    const utils::WgCount wgc = utils::compute_2d_workgroup_count(
+        g.device(), om.numel, wg_size, "slice(resize)");
+    set_slice_dispatch_grid(g, dispatch_idx, wgc);
   };
   if (is_symint(graph, start_id)) {
     graph.add_resize_hook(start_id, recompute);
@@ -244,6 +405,20 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   graph.own_uniform_buffer(out_meta_buf);
   graph.own_uniform_buffer(in_meta_buf);
   graph.own_uniform_buffer(params_buf);
+
+  // Offer this dispatch to a following whole-extent copy of `out`.
+  WebGPUGraph::SliceChain chain;
+  chain.valid = (wg_size == kSliceDualWorkgroupSizeX);
+  chain.out_id = out_id;
+  chain.dispatch_idx = dispatch_idx;
+  chain.in_buffer = in_tensor.buffer;
+  chain.in_nbytes = in_tensor.nbytes;
+  chain.out_buffer = out_tensor.buffer;
+  chain.out_nbytes = out_tensor.nbytes;
+  chain.out_meta_buf = out_meta_buf;
+  chain.in_meta_buf = in_meta_buf;
+  chain.params_buf = params_buf;
+  graph.offer_slice_chain(chain);
 }
 
 } // namespace
diff --git a/backends/webgpu/runtime/ops/slice/SliceDispatch.h b/backends/webgpu/runtime/ops/slice/SliceDispatch.h
new file mode 100644
index 00000000000..a3da993891f
--- /dev/null
+++ b/backends/webgpu/runtime/ops/slice/SliceDispatch.h
@@ -0,0 +1,27 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/backends/webgpu/runtime/WebGPUDispatchMath.h>
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+
+#include <cstddef>
+
+namespace executorch::backends::webgpu {
+
+inline void set_slice_dispatch_grid(
+    WebGPUGraph& graph,
+    size_t dispatch_index,
+    const utils::WgCount& grid) {
+  WebGPUDispatch& dispatch = graph.dispatch_at(dispatch_index);
+  dispatch.workgroup_count_x = grid.x;
+  dispatch.workgroup_count_y = grid.y;
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/slice/slice.wgsl b/backends/webgpu/runtime/ops/slice/slice.wgsl
index e6940ee2e4e..7fa06e5897c 100644
--- a/backends/webgpu/runtime/ops/slice/slice.wgsl
+++ b/backends/webgpu/runtime/ops/slice/slice.wgsl
@@ -20,8 +20,11 @@ struct Params {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let out_bufi = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-spill: numel can exceed the 65535 per-dim grid cap.
+    let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (out_bufi >= out_meta.numel) {
         return;
     }
diff --git a/backends/webgpu/runtime/ops/slice/slice_dual.wgsl b/backends/webgpu/runtime/ops/slice/slice_dual.wgsl
new file mode 100644
index 00000000000..2ef4dc92d74
--- /dev/null
+++ b/backends/webgpu/runtime/ops/slice/slice_dual.wgsl
@@ -0,0 +1,65 @@
+// 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> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: array<vec4<u32>, 2>,
+  strides: array<vec4<u32>, 2>,
+}
+@group(0) @binding(2) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(3) var<uniform> in_meta: TensorMeta;
+
+struct Params {
+  dim: u32,
+  start: u32,
+  step: u32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+// Second destination for a chained aten.slice_copy.Tensor whose range is the
+// whole of `output` (start == 0, step == 1, end >= size on its dim), i.e. a
+// pure elementwise copy of `output`. Same flat index, so one gather feeds both
+// stores. `output` is still written, so no consumer of it is assumed dead.
+@group(0) @binding(5) var<storage, read_write> output2: array<f32>;
+
+// Fixed workgroup size (Apple scalar ALU: no override loop bounds). 64 is the
+// value clamp_workgroup_size already yields for slice.wgsl.
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-spill: numel can exceed the 65535 per-dim grid cap. Stride is the
+    // literal 64 above, which is the real threads-per-group either way.
+    let out_bufi = gid.x + gid.y * (num_workgroups.x * 64u);
+    if (out_bufi >= out_meta.numel) {
+        return;
+    }
+
+    // Gather: out_bufi -> in_bufi, sliced dim coord = start + coord*step.
+    var rem = out_bufi;
+    var in_bufi: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d >> 2u][d & 3u];
+        rem = rem % out_meta.strides[d >> 2u][d & 3u];
+        var in_coord = coord;
+        if (d == params.dim) {
+            in_coord = params.start + coord * params.step;
+        }
+        in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u];
+    }
+    let v = input[in_bufi];
+    output[out_bufi] = v;
+    output2[out_bufi] = v;
+}
+// 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.
diff --git a/backends/webgpu/runtime/ops/slice/slice_dual_guard.h b/backends/webgpu/runtime/ops/slice/slice_dual_guard.h
new file mode 100644
index 00000000000..55ce658e933
--- /dev/null
+++ b/backends/webgpu/runtime/ops/slice/slice_dual_guard.h
@@ -0,0 +1,57 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// Pure predicate for the slice dual-store dispatch merge; unit-tested by
+// slice_dual_guard_test so the load-bearing condition is checked without a GPU.
+//
+// A chained aten.slice_copy.Tensor is a PURE ELEMENTWISE COPY of its input --
+// out2[i] == out1[i] at the identical flat index -- only when it starts at a
+// STATIC 0, steps by 1, and its STATIC end already covers the serialized
+// maximum extent of the sliced dim. norm_clamp() then lands on the live size at
+// every live shape, so the identity holds for the whole dynamic-shape range.
+//
+// A SymInt bound is rejected: it can resolve below the live extent at execute
+// time, at which point the second slice is a strict subset and copying the
+// whole gather into it would be wrong.
+struct SliceDualSpan {
+  int64_t step = 0;
+  bool start_is_symint = true;
+  bool end_is_symint = true;
+  int64_t start = 0; // resolved at the serialized maximum shape
+  int64_t end = 0; // resolved at the serialized maximum shape
+  int64_t dim_size = 0; // serialized maximum extent of the sliced dim
+};
+
+inline bool slice_dual_full_span(const SliceDualSpan& s) {
+  if (s.step != 1) {
+    return false;
+  }
+  if (s.start_is_symint || s.end_is_symint) {
+    return false;
+  }
+  if (s.dim_size <= 0) {
+    return false;
+  }
+  return s.start == 0 && s.end >= s.dim_size;
+}
+
+// The fused dispatch gathers from `in` and stores to `out1` and `out2` in one
+// invocation. `out2` must therefore not be either of the buffers the gather
+// already reads or writes.
+inline bool
+slice_dual_buffers_ok(const void* out2, const void* out1, const void* in) {
+  return out2 != nullptr && out2 != out1 && out2 != in;
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/slice/slice_dual_wgsl.h b/backends/webgpu/runtime/ops/slice/slice_dual_wgsl.h
new file mode 100644
index 00000000000..972818d321f
--- /dev/null
+++ b/backends/webgpu/runtime/ops/slice/slice_dual_wgsl.h
@@ -0,0 +1,89 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from slice_dual.wgsl - DO NOT EDIT.
+// wgsl-sha256: 845ec79985822d5621fd16688a45af8f363310f949ea4ac14c77bf0e84f9c9c0
+inline constexpr const char* kSliceDualWGSL = R"(
+// 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> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: array<vec4<u32>, 2>,
+  strides: array<vec4<u32>, 2>,
+}
+@group(0) @binding(2) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(3) var<uniform> in_meta: TensorMeta;
+
+struct Params {
+  dim: u32,
+  start: u32,
+  step: u32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+// Second destination for a chained aten.slice_copy.Tensor whose range is the
+// whole of `output` (start == 0, step == 1, end >= size on its dim), i.e. a
+// pure elementwise copy of `output`. Same flat index, so one gather feeds both
+// stores. `output` is still written, so no consumer of it is assumed dead.
+@group(0) @binding(5) var<storage, read_write> output2: array<f32>;
+
+// Fixed workgroup size (Apple scalar ALU: no override loop bounds). 64 is the
+// value clamp_workgroup_size already yields for slice.wgsl.
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-spill: numel can exceed the 65535 per-dim grid cap. Stride is the
+    // literal 64 above, which is the real threads-per-group either way.
+    let out_bufi = gid.x + gid.y * (num_workgroups.x * 64u);
+    if (out_bufi >= out_meta.numel) {
+        return;
+    }
+
+    // Gather: out_bufi -> in_bufi, sliced dim coord = start + coord*step.
+    var rem = out_bufi;
+    var in_bufi: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d >> 2u][d & 3u];
+        rem = rem % out_meta.strides[d >> 2u][d & 3u];
+        var in_coord = coord;
+        if (d == params.dim) {
+            in_coord = params.start + coord * params.step;
+        }
+        in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u];
+    }
+    let v = input[in_bufi];
+    output[out_bufi] = v;
+    output2[out_bufi] = v;
+}
+// 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.
+)";
+
+inline constexpr uint32_t kSliceDualWorkgroupSizeX = 64;
+inline constexpr uint32_t kSliceDualWorkgroupSizeY = 1;
+inline constexpr uint32_t kSliceDualWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/slice/slice_wgsl.h b/backends/webgpu/runtime/ops/slice/slice_wgsl.h
index ff70bacd4a6..783f7cfc6a3 100644
--- a/backends/webgpu/runtime/ops/slice/slice_wgsl.h
+++ b/backends/webgpu/runtime/ops/slice/slice_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from slice.wgsl - DO NOT EDIT.
-// wgsl-sha256: 6a895a8c321cd3ddaffc468d7843c9dea3eaeb9d3de0088a1a09419b3ccfd10b
+// wgsl-sha256: 5b9cb0c437a87fac0e66c1a6e7fedba75068737e2b9ed03f53b4f6ec1250081d
 inline constexpr const char* kSliceWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<f32>;
@@ -37,8 +37,11 @@ struct Params {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let out_bufi = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-spill: numel can exceed the 65535 per-dim grid cap.
+    let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (out_bufi >= out_meta.numel) {
         return;
     }
diff --git a/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp
index 5bdba562633..eac36e61765 100644
--- a/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp
+++ b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp
@@ -9,28 +9,77 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/update_cache/UpdateCacheState.h>
 #include <executorch/backends/webgpu/runtime/ops/update_cache/update_cache_wgsl.h>
 
 #include <webgpu/webgpu.h>
 
 #include <cstdint>
-#include <cstring>
+#include <limits>
 #include <stdexcept>
 
 namespace executorch::backends::webgpu {
 
 namespace {
 
-// Uniform buffer layout matching the WGSL Params struct (16-byte aligned).
-struct UpdateCacheParams {
-  uint32_t numel;
-  uint32_t dst_offset;
-  uint32_t cache_numel;
-  uint32_t _pad0;
+int64_t read_input_pos(const WebGPUGraph& graph, int input_pos_id) {
+  const auto type = graph.get_value_type(input_pos_id);
+  if (type == WebGPUGraph::ValueType::Int) {
+    return graph.get_int(input_pos_id);
+  }
+  if (type == WebGPUGraph::ValueType::SymInt) {
+    return graph.read_symint(input_pos_id);
+  }
+  throw std::runtime_error(
+      "WebGPU update_cache: input_pos must be Int or SymInt");
+}
+
+void validate_fp32_tensor(const WebGPUTensor& tensor, const char* label) {
+  const uint64_t numel = utils::numel_of(tensor.dims);
+  if (numel > std::numeric_limits<size_t>::max() / sizeof(float) ||
+      tensor.nbytes != static_cast<size_t>(numel) * sizeof(float)) {
+    throw std::runtime_error(
+        std::string("WebGPU update_cache: ") + label + " must be fp32");
+  }
+}
+
+struct UpdateCacheRefreshContext {
+  int value_id;
+  int cache_id;
+  int input_pos_id;
+  size_t expected_value_rank;
+  size_t expected_cache_rank;
+  uint32_t workgroup_size;
+  uint32_t max_workgroups_per_dimension;
+  WGPUBuffer params_buffer;
+  size_t dispatch_index;
 };
-static_assert(
-    sizeof(UpdateCacheParams) == 16,
-    "UpdateCacheParams must be 16 bytes");
+
+void refresh_update_cache(
+    WebGPUGraph& graph,
+    const UpdateCacheRefreshContext& context) {
+  const LiveUpdateCacheInputs inputs = {
+      graph.cur_dims(context.value_id),
+      graph.cur_dims(context.cache_id),
+      context.expected_value_rank,
+      context.expected_cache_rank,
+      read_input_pos(graph, context.input_pos_id),
+      context.workgroup_size,
+      context.max_workgroups_per_dimension,
+  };
+  refresh_live_update_cache_state(
+      inputs, [&](const LiveUpdateCacheState& state) {
+        wgpuQueueWriteBuffer(
+            graph.queue(),
+            context.params_buffer,
+            0,
+            &state.params,
+            sizeof(state.params));
+        auto& dispatch = graph.dispatch_at(context.dispatch_index);
+        dispatch.workgroup_count_x = state.workgroup_count_x;
+        dispatch.workgroup_count_y = 1;
+      });
+}
 
 // llama.update_cache.default args: [value, cache, input_pos, out].
 void update_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
@@ -42,82 +91,29 @@ void update_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
   const auto& value_tensor = graph.get_tensor(value_id);
   const auto& cache_tensor = graph.get_tensor(cache_id);
-  if (value_tensor.dims.size() < 4 || cache_tensor.dims.size() < 4 ||
-      value_tensor.nbytes == 0) {
+  if (value_tensor.dims.size() < 4 || cache_tensor.dims.size() < 4) {
     throw std::runtime_error("WebGPU update_cache: expects 4D value and cache");
   }
-
-  uint64_t value_numel = 1;
-  for (int64_t d : value_tensor.dims) {
-    value_numel *= static_cast<uint64_t>(d);
-  }
-  // fp32-only shader: bail if bytes don't match an fp32 element count.
-  if (value_tensor.nbytes != value_numel * sizeof(float)) {
-    throw std::runtime_error(
-        "WebGPU update_cache: fp32-only (byte-size mismatch)");
-  }
-
-  const size_t ndim = value_tensor.dims.size();
-  const size_t cndim = cache_tensor.dims.size();
-  // Mirror Vulkan update_cache_impl shape guards (backends/vulkan SDPA.cpp).
-  if (value_tensor.dims[ndim - 4] != 1 || cache_tensor.dims[cndim - 4] != 1) {
-    throw std::runtime_error("WebGPU update_cache: batch must be 1");
-  }
-  if (value_tensor.dims[ndim - 1] != cache_tensor.dims[cndim - 1]) {
-    throw std::runtime_error("WebGPU update_cache: head_dim mismatch");
-  }
-  if (value_tensor.dims[ndim - 2] != cache_tensor.dims[cndim - 2]) {
-    throw std::runtime_error("WebGPU update_cache: n_heads mismatch");
-  }
-  const uint64_t head_dim = static_cast<uint64_t>(value_tensor.dims[ndim - 1]);
-  const uint64_t n_heads = static_cast<uint64_t>(value_tensor.dims[ndim - 2]);
-
-  uint64_t cache_numel = 1;
-  for (int64_t d : cache_tensor.dims) {
-    cache_numel *= static_cast<uint64_t>(d);
-  }
-
-  if (graph.get_value_type(input_pos_id) != WebGPUGraph::ValueType::Int) {
-    throw std::runtime_error(
-        "WebGPU update_cache: input_pos must be Int (SymInt not yet supported)");
-  }
-  const int64_t input_pos = graph.get_int(input_pos_id);
-  if (input_pos < 0) {
-    throw std::runtime_error(
-        "WebGPU update_cache: input_pos must be non-negative");
-  }
-
-  // Bound input_pos in u64 so the u32 param downcasts cannot overflow/truncate.
-  const uint64_t stride = n_heads * head_dim;
-  if (cache_numel > UINT32_MAX || value_numel > cache_numel ||
-      static_cast<uint64_t>(input_pos) > (cache_numel - value_numel) / stride) {
-    throw std::runtime_error(
-        "WebGPU update_cache: input_pos writes past cache capacity");
-  }
-  const uint64_t dst_offset = static_cast<uint64_t>(input_pos) * stride;
-
-  UpdateCacheParams params = {};
-  params.numel = static_cast<uint32_t>(value_numel);
-  params.dst_offset = static_cast<uint32_t>(dst_offset);
-  params.cache_numel = static_cast<uint32_t>(cache_numel);
+  validate_fp32_tensor(value_tensor, "value");
+  validate_fp32_tensor(cache_tensor, "cache");
 
   // Validate dispatch against device limits before allocating GPU objects.
   const uint32_t wg_size =
       utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX);
-  const uint32_t workgroup_count_x = utils::compute_1d_workgroup_count(
-      device, params.numel, wg_size, "update_cache");
-
-  WGPUBufferDescriptor uniform_desc = {};
-  uniform_desc.size = sizeof(UpdateCacheParams);
-  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
-  uniform_desc.mappedAtCreation = true;
-  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);
-  void* mapped =
-      wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(UpdateCacheParams));
-  std::memcpy(mapped, &params, sizeof(UpdateCacheParams));
-  wgpuBufferUnmap(uniform_buffer);
-
-  graph.add_uniform_buffer_bytes(sizeof(UpdateCacheParams));
+  const uint32_t max_workgroups = utils::queried_max_workgroups(device);
+  const LiveUpdateCacheInputs initial_inputs = {
+      value_tensor.dims,
+      cache_tensor.dims,
+      value_tensor.dims.size(),
+      cache_tensor.dims.size(),
+      read_input_pos(graph, input_pos_id),
+      wg_size,
+      max_workgroups,
+  };
+  const LiveUpdateCacheState initial_state =
+      compute_live_update_cache_state(initial_inputs);
+  WGPUBuffer uniform_buffer =
+      graph.create_params_buffer(initial_state.params);
 
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
@@ -143,10 +139,32 @@ void update_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       &wg_size_constant,
       1);
 
-  graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count_x});
-
-  // Drop our ref; the bind group keeps the uniform buffer alive until release.
-  wgpuBufferRelease(uniform_buffer);
+  const size_t dispatch_index = graph.add_dispatch(
+      {bundle.pipeline,
+       bundle.bind_group,
+       initial_state.workgroup_count_x,
+       "update_cache"});
+
+  const UpdateCacheRefreshContext refresh_context = {
+      value_id,
+      cache_id,
+      input_pos_id,
+      value_tensor.dims.size(),
+      cache_tensor.dims.size(),
+      wg_size,
+      max_workgroups,
+      uniform_buffer,
+      dispatch_index,
+  };
+  std::vector<int> symint_triggers;
+  if (graph.get_value_type(input_pos_id) == WebGPUGraph::ValueType::SymInt) {
+    symint_triggers.push_back(input_pos_id);
+  }
+  graph.add_post_resize_hook(
+      {value_id, cache_id},
+      symint_triggers,
+      refresh_update_cache,
+      refresh_context);
 }
 
 } // namespace
diff --git a/backends/webgpu/runtime/ops/update_cache/UpdateCacheState.h b/backends/webgpu/runtime/ops/update_cache/UpdateCacheState.h
new file mode 100644
index 00000000000..234b5185ec4
--- /dev/null
+++ b/backends/webgpu/runtime/ops/update_cache/UpdateCacheState.h
@@ -0,0 +1,141 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/backends/webgpu/runtime/WebGPUDispatchMath.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <stdexcept>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+struct UpdateCacheParams {
+  uint32_t numel;
+  uint32_t dst_offset;
+  uint32_t cache_numel;
+  uint32_t _pad0;
+};
+static_assert(
+    sizeof(UpdateCacheParams) == 16,
+    "UpdateCacheParams must be 16 bytes");
+
+struct LiveUpdateCacheInputs {
+  std::vector<int64_t> value_dims;
+  std::vector<int64_t> cache_dims;
+  size_t expected_value_rank;
+  size_t expected_cache_rank;
+  int64_t start_pos;
+  uint32_t workgroup_size;
+  uint32_t max_workgroups_per_dimension;
+};
+
+struct LiveUpdateCacheState {
+  UpdateCacheParams params;
+  uint32_t workgroup_count_x;
+};
+
+inline LiveUpdateCacheState compute_live_update_cache_state(
+    const LiveUpdateCacheInputs& inputs) {
+  if (inputs.value_dims.size() != inputs.expected_value_rank ||
+      inputs.cache_dims.size() != inputs.expected_cache_rank ||
+      inputs.value_dims.size() < 4 || inputs.cache_dims.size() < 4) {
+    throw std::runtime_error("WebGPU update_cache: tensor rank changed");
+  }
+  for (int64_t dim : inputs.value_dims) {
+    if (dim <= 0) {
+      throw std::runtime_error(
+          "WebGPU update_cache: value dimensions must be positive");
+    }
+  }
+  for (int64_t dim : inputs.cache_dims) {
+    if (dim <= 0) {
+      throw std::runtime_error(
+          "WebGPU update_cache: cache dimensions must be positive");
+    }
+  }
+
+  const size_t value_rank = inputs.value_dims.size();
+  const size_t cache_rank = inputs.cache_dims.size();
+  if (inputs.value_dims[value_rank - 4] != 1 ||
+      inputs.cache_dims[cache_rank - 4] != 1) {
+    throw std::runtime_error("WebGPU update_cache: batch must be 1");
+  }
+  if (inputs.value_dims[value_rank - 2] !=
+      inputs.cache_dims[cache_rank - 2]) {
+    throw std::runtime_error("WebGPU update_cache: n_heads mismatch");
+  }
+  if (inputs.value_dims[value_rank - 1] !=
+      inputs.cache_dims[cache_rank - 1]) {
+    throw std::runtime_error("WebGPU update_cache: head_dim mismatch");
+  }
+
+  const uint64_t value_numel = utils::numel(inputs.value_dims);
+  const uint64_t cache_numel = utils::numel(inputs.cache_dims);
+  const uint64_t heads =
+      static_cast<uint64_t>(inputs.value_dims[value_rank - 2]);
+  const uint64_t head_dim =
+      static_cast<uint64_t>(inputs.value_dims[value_rank - 1]);
+  if (heads > std::numeric_limits<uint64_t>::max() / head_dim) {
+    throw std::runtime_error("WebGPU update_cache: stride overflow");
+  }
+  const uint64_t stride = heads * head_dim;
+  if (stride == 0) {
+    throw std::runtime_error("WebGPU update_cache: stride must be positive");
+  }
+  if (inputs.start_pos < 0) {
+    throw std::runtime_error(
+        "WebGPU update_cache: input_pos must be non-negative");
+  }
+  const uint64_t start_pos = static_cast<uint64_t>(inputs.start_pos);
+  if (start_pos > std::numeric_limits<uint64_t>::max() / stride) {
+    throw std::runtime_error("WebGPU update_cache: input_pos offset overflow");
+  }
+  const uint64_t dst_offset = start_pos * stride;
+
+  constexpr uint64_t kMaxU32 = std::numeric_limits<uint32_t>::max();
+  if (cache_numel > kMaxU32 || value_numel > cache_numel ||
+      value_numel > kMaxU32 || dst_offset > kMaxU32 ||
+      dst_offset > cache_numel - value_numel) {
+    throw std::runtime_error(
+        "WebGPU update_cache: input_pos writes past cache capacity");
+  }
+  if (inputs.workgroup_size == 0 ||
+      inputs.max_workgroups_per_dimension == 0) {
+    throw std::runtime_error(
+        "WebGPU update_cache: dispatch limits must be positive");
+  }
+  const uint64_t workgroup_count = value_numel / inputs.workgroup_size +
+      static_cast<uint64_t>(value_numel % inputs.workgroup_size != 0);
+  if (workgroup_count == 0 ||
+      workgroup_count > inputs.max_workgroups_per_dimension) {
+    throw std::runtime_error(
+        "WebGPU update_cache: workgroup count exceeds the 1D dispatch limit");
+  }
+
+  LiveUpdateCacheState state = {};
+  state.params.numel = static_cast<uint32_t>(value_numel);
+  state.params.dst_offset = static_cast<uint32_t>(dst_offset);
+  state.params.cache_numel = static_cast<uint32_t>(cache_numel);
+  state.workgroup_count_x = static_cast<uint32_t>(workgroup_count);
+  return state;
+}
+
+template <typename Commit>
+void refresh_live_update_cache_state(
+    const LiveUpdateCacheInputs& inputs,
+    Commit&& commit) {
+  const LiveUpdateCacheState state =
+      compute_live_update_cache_state(inputs);
+  commit(state);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py
index 98812c0f144..f2c750a50ac 100644
--- a/backends/webgpu/test/test_wgsl_codegen.py
+++ b/backends/webgpu/test/test_wgsl_codegen.py
@@ -240,14 +240,14 @@ def test_generated_output_manifest_digest(self) -> None:
             digest.update(b"\0")
             digest.update(output.read_bytes())
             digest.update(b"\0")
-        self.assertEqual(len(outputs), 138)
+        self.assertEqual(len(outputs), 144)
         self.assertEqual(
             digest.hexdigest(),
-            "fee848cd069b4c09d3d2e9a7920331f46d5646b74bec30259542dde8f287e504",
+            "b4a5a79ea7cdd1f18867106365da79b6faaa5f1065409cadec8bab2d64f3139e",
         )
         self.assertEqual(
             hashlib.sha256(g.registry_path().read_bytes()).hexdigest(),
-            "477721998b3cd8f3f0fdd485fa797c71035a20cbc10a8b4bf44893e37fa435b8",
+            "2717f916578362e00727dabd1d7e91a28c8cada0c238fa6dc323511cd672369a",
         )
 
     def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: