From 5ae9624ce1c96e7b9691970aade7770fdfd538fb Mon Sep 17 00:00:00 2001
From: Siddartha Pothapragada <sidart@meta.com>
Date: Tue, 20 Oct 2026 22:15:38 -0700
Subject: [PATCH] Fix the fused QKV resize check for dual-route projections

Summary:
On GPUs with `shader-f16`, Llama 3.2 1B's Q/K/V projections are fused by `QkvBk64`. When the live M doesn't use the fused kernel (any M other than 128, 508 or 512, including decode), the fusion's resize hook checks that the separate projections restored their dispatches. It required every dispatch in each member's range to be live. With a dynamic M each `linear_q4gsw` member records a decode and a prefill route, and `select_dispatch_route` keeps the inactive route at a zero grid, so the check threw on every unfused call: "linear_q4gsw_bk64_qkv(resize): separate projection dispatch was not restored before the QKV resize hook ran".

This breaks Llama 3.2 1B chat in the ExecuTorch WebGPU Playground on executorch.ai on Apple GPUs (reported on Metal 3). The check now requires one live dispatch per member, which still catches a member whose hook didn't run.

CI runs on SwiftShader, which has no `shader-f16`, so the fused path and `DynamicShape.QkvBk64LiveRoutesProfile` (M = 128, 1, 128) don't run there.

Differential Revision: D126153972
---
 backends/webgpu/runtime/passes/QkvBk64.cpp | 26 +++++++++++++---------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/backends/webgpu/runtime/passes/QkvBk64.cpp b/backends/webgpu/runtime/passes/QkvBk64.cpp
index 336dbbc446f..9022ae437f4 100644
--- a/backends/webgpu/runtime/passes/QkvBk64.cpp
+++ b/backends/webgpu/runtime/passes/QkvBk64.cpp
@@ -110,21 +110,25 @@ void resize_qkv_bk64(WebGPUGraph& graph, const QkvBk64ResizeContext& context) {
   } else {
     // The separate projections' own resize hooks are registered before this
     // one (Phase 3 processes each Q/K/V member before this fusion's combined
-    // hook) and unconditionally restore their live grids, so this hook
-    // normally has nothing to do here. That ordering isn't enforced by the
-    // type system, so fail loud rather than silently drop Q/K/V outputs if a
-    // future change ever violates it.
+    // hook) and restore their live grids, so this hook normally has nothing to
+    // do here. A member that records a decode and a prefill route keeps its
+    // inactive route at a zero grid, so require one live dispatch per member
+    // rather than all of them. That ordering isn't enforced by the type
+    // system, so fail loud rather than silently drop Q/K/V outputs if a future
+    // change ever violates it.
     for (size_t member = 0; member < context.separate_begin.size(); member++) {
+      bool live = false;
       for (size_t i = context.separate_begin[member];
-           i < context.separate_end[member];
+           i < context.separate_end[member] && !live;
            i++) {
         const auto& dispatch = graph.dispatch_at(i);
-        if (dispatch.workgroup_count_x == 0u ||
-            dispatch.workgroup_count_y == 0u) {
-          throw std::runtime_error(
-              "linear_q4gsw_bk64_qkv(resize): separate projection dispatch "
-              "was not restored before the QKV resize hook ran");
-        }
+        live = dispatch.workgroup_count_x != 0u &&
+            dispatch.workgroup_count_y != 0u;
+      }
+      if (!live) {
+        throw std::runtime_error(
+            "linear_q4gsw_bk64_qkv(resize): separate projection dispatch "
+            "was not restored before the QKV resize hook ran");
       }
     }
   }