[ET-VK][sdpa] Vendor-adaptive head_dim output-tiling (TILE_N4=2) in GQA AV coop-GEMV

Differential Revision: D112906313

Pull Request resolved: #21064

Author: SS-JIA

backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml

@@ -26,3 +26,11 @@ sdpa_compute_out_coop:
     - NAME: sdpa_compute_out_coop
     - NAME: sdpa_compute_out_gqa_coop
       GQA: True
+    # head_dim output-tiled GQA variant: each workgroup owns TILE_N4 head_dim
+    # texels, amortizing the attn-weight loads and shared-memory reduction over
+    # G x TILE_N4 outputs. A win on Adreno but a regression on Mali at common
+    # decode contexts, so it is selected vendor-adaptively (Adreno only) in
+    # SDPA.cpp; the partial_n_tile guard handles D4 % TILE_N4 != 0.
+    - NAME: sdpa_compute_out_gqa_coop_tile2
+      GQA: True
+      TILE_N4: 2

backends/vulkan/runtime/graph/ops/impl/SDPA.cpp

@@ -190,8 +190,9 @@ bool use_gqa_av_coop(
 }
 
 // Resolve whether the AV decode path uses the GQA-reuse coop shader, honoring
-// the test-only shader_override knob (see SDPA.h): -1/kDummyValueRef
-// auto-selects via use_gqa_av_coop; 0 forces off; 1 forces on.
+// the test-only shader_override knob (see SDPA.h): auto selects via
+// use_gqa_av_coop; kShaderOverrideForceNonGqa forces off; any other value
+// forces the GQA family on.
 bool resolve_use_gqa(
     ComputeGraph* graph,
     const ValueRef shader_override,
@@ -200,20 +201,38 @@ bool resolve_use_gqa(
   if (shader_override == kDummyValueRef) {
     return use_gqa_av_coop(graph, num_q_heads, num_kv_heads);
   }
-  const bool force_gqa = graph->extract_scalar<int64_t>(shader_override) != 0;
+  const bool force_gqa = graph->extract_scalar<int64_t>(shader_override) !=
+      kShaderOverrideForceNonGqa;
   // Forcing the GQA shader on an ineligible shape would silently drop query
   // heads (z-dispatch = Hkv cannot cover Hq, or group size exceeds the shader's
   // fixed accumulator array). Fail loudly instead of producing garbage output.
   if (force_gqa) {
     VK_CHECK_COND(
         num_kv_heads > 0 && num_q_heads % num_kv_heads == 0 &&
             num_q_heads / num_kv_heads <= kMaxGqaGroupSize,
-        "shader_override=1 requires a GQA-eligible shape: Hq divisible by Hkv and "
-        "group size <= kMaxGqaGroupSize");
+        "forcing GQA via shader_override requires a GQA-eligible shape: Hq "
+        "divisible by Hkv and group size <= kMaxGqaGroupSize");
   }
   return force_gqa;
 }
 
+// Resolve whether the GQA-reuse AV path uses the head_dim output-tiled (_tile2)
+// variant. Auto (no override) and kShaderOverrideForceGqa defer to the vendor
+// gate (tile2 on Adreno); the force values pin the variant on any device (see
+// SDPA.h), giving the Adreno-only tile2 variant deterministic test coverage.
+bool resolve_use_tile2(ComputeGraph* graph, const ValueRef shader_override) {
+  if (shader_override != kDummyValueRef) {
+    const int64_t ov = graph->extract_scalar<int64_t>(shader_override);
+    if (ov == kShaderOverrideForceTile2) {
+      return true;
+    }
+    if (ov == kShaderOverrideForceBase) {
+      return false;
+    }
+  }
+  return graph->device_is_adreno();
+}
+
 vkapi::ShaderInfo pick_sdpa_qk_shader(
     ComputeGraph* graph,
     const std::vector<ArgGroup>& args,
@@ -347,6 +366,16 @@ vkapi::ShaderInfo pick_sdpa_av_shader(
         // reusing it across the group. Cuts the dominant V-cache traffic ~Gx
         // for this bandwidth-bound kernel.
         shader_name += "_gqa_coop";
+        // The _tile2 (head_dim output-tiled) variant amortizes the attn-weight
+        // loads and shared-memory reduction over more outputs per workgroup. A
+        // consistent win on Adreno (AV ~1.14-1.63x) but a regression on Mali at
+        // common decode contexts (~0.67-0.86x, interleaved median), so it is
+        // selected vendor-adaptively (tests can pin it via shader_override; see
+        // resolve_use_tile2). pick_sdpa_av_global_wg_size keys the x-dim
+        // collapse off this same _tile2 suffix.
+        if (resolve_use_tile2(graph, shader_override)) {
+          shader_name += "_tile2";
+        }
       } else {
         shader_name += "_coop";
       }
@@ -389,7 +418,14 @@ utils::uvec3 pick_sdpa_av_global_wg_size(
   if (shader.kernel_name.find("_gqa_coop") != std::string::npos) {
     const ValueRef v = args.at(1).refs.at(1);
     const uint32_t num_kv_heads = graph->size_at<uint32_t>(-2, v);
-    return {N4, M4, static_cast<uint32_t>(num_kv_heads * d.B)};
+    // The _tile2 variant has each workgroup own 2 head_dim texels, so its x-dim
+    // collapses to div_up(D4, 2). The plain (TILE_N4 == 1) variant keeps x ==
+    // D4.
+    const uint32_t x_dim =
+        shader.kernel_name.find("_tile2") != std::string::npos
+        ? utils::div_up(N4, 2u)
+        : N4;
+    return {x_dim, M4, static_cast<uint32_t>(num_kv_heads * d.B)};
   }
 
   return {N4, M4, static_cast<uint32_t>(d.H * d.B)};

backends/vulkan/runtime/graph/ops/impl/SDPA.h

@@ -44,16 +44,31 @@ void add_sdpa_attn_weights_softmax_node(
     const ValueRef attn_weights_softmax,
     const SDPAMode mode);
 
+// Scalar values for the shader_override knob (see add_sdpa_compute_out_node).
+constexpr int64_t kShaderOverrideForceNonGqa = 0;
+constexpr int64_t kShaderOverrideForceGqa = 1;
+constexpr int64_t kShaderOverrideForceTile2 = 2;
+constexpr int64_t kShaderOverrideForceBase = 3;
+
 // shader_override: test-only knob selecting the single-token decode AV shader.
-//   kDummyValueRef (-1) — auto: use the GQA-reuse coop shader
-//                         (`sdpa_compute_out_gqa_coop`) when it applies.
-//   0                   — force the per-query-head coop shader
-//                         (`sdpa_compute_out_coop`).
-//   1                   — force the GQA-reuse coop shader.
-// Lets the benchmark/test exercise both AV shaders on the same shape. Forcing
-// GQA (1) requires a GQA-eligible shape (Hq divisible by Hkv, group size <= the
-// shader's compile-time bound); this is VK_CHECK'd, so an ineligible forced
-// shape fails loudly rather than silently dropping query heads.
+// It is a ValueRef holding one of the kShaderOverride* int64 scalars above, or
+// kDummyValueRef for "no override" (auto).
+//   auto (kDummyValueRef)   — use the GQA-reuse coop shader when it applies,
+//                             with the tiled vs base variant chosen by vendor.
+//   kShaderOverrideForceNonGqa — force the per-query-head coop shader
+//                             (`sdpa_compute_out_coop`).
+//   kShaderOverrideForceGqa    — force the GQA-reuse coop shader; vendor picks
+//   the
+//                             base (`sdpa_compute_out_gqa_coop`) vs head_dim
+//                             output-tiled (`sdpa_compute_out_gqa_coop_tile2`)
+//                             variant (tile2 on Adreno).
+//   kShaderOverrideForceTile2  — force the tiled variant regardless of vendor.
+//   kShaderOverrideForceBase   — force the base variant regardless of vendor.
+// Lets the benchmark/test exercise every AV shader on the same shape, and gives
+// the Adreno-only tile2 variant deterministic coverage on any device. Forcing
+// the GQA family requires a GQA-eligible shape (Hq divisible by Hkv, group size
+// <= the shader's compile-time bound); this is VK_CHECK'd, so an ineligible
+// forced shape fails loudly rather than silently dropping query heads.
 void add_sdpa_compute_out_node(
     ComputeGraph& graph,
     const ValueRef attn_weights_softmax,

backends/vulkan/test/custom_ops/impl/TestSDPA.cpp

@@ -119,9 +119,11 @@ static void test_sdpa_impl(
 // Decode vs prefill is selected automatically inside the nodes via
 // is_single_token() (S == 1 -> coop/GEMV shaders, S > 1 -> tiled shaders). For
 // single-token decode the AV shader is picked by impl_selector:
-//   "default" -> auto-select (GQA-reuse coop shader when it applies)
-//   "gqa"     -> force the GQA-reuse coop shader
-//   "non_gqa" -> force the per-query-head coop shader
+//   "default"   -> auto-select (GQA-reuse coop shader when it applies)
+//   "gqa"       -> force the GQA-reuse coop shader (vendor picks base vs tile2)
+//   "gqa_tile2" -> force the head_dim output-tiled GQA variant (any device)
+//   "gqa_base"  -> force the base (non-tiled) GQA variant (any device)
+//   "non_gqa"   -> force the per-query-head coop shader
 // impl_selector has no effect on prefill (tiled).
 //
 // Args: q, k_cache, v_cache, impl_selector, out
@@ -136,8 +138,10 @@ void test_sdpa(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   const std::string impl_selector = graph.extract_string(impl_selector_str);
   VK_CHECK_COND(
       impl_selector == "default" || impl_selector == "gqa" ||
+          impl_selector == "gqa_tile2" || impl_selector == "gqa_base" ||
           impl_selector == "non_gqa",
-      "test_sdpa: impl_selector must be one of {default, gqa, non_gqa}");
+      "test_sdpa: impl_selector must be one of {default, gqa, gqa_tile2, "
+      "gqa_base, non_gqa}");
 
   const int64_t seq_len = graph.size_at<int64_t>(-3, q);
   const int64_t context_len = graph.size_at<int64_t>(-3, k_cache);
@@ -147,12 +151,17 @@ void test_sdpa(ComputeGraph& graph, const std::vector<ValueRef>& args) {
 
   const ValueRef input_pos_symint = graph.add_symint(input_pos_val);
 
-  // shader_override (see SDPA.h): -1 auto, 0 force non-GQA, 1 force GQA.
+  // shader_override (see SDPA.h): kDummyValueRef auto; otherwise a
+  // kShaderOverride* scalar forcing the AV shader family / variant.
   ValueRef shader_override = kDummyValueRef;
   if (impl_selector == "gqa") {
-    shader_override = graph.add_scalar<int64_t>(1);
+    shader_override = graph.add_scalar<int64_t>(kShaderOverrideForceGqa);
+  } else if (impl_selector == "gqa_tile2") {
+    shader_override = graph.add_scalar<int64_t>(kShaderOverrideForceTile2);
+  } else if (impl_selector == "gqa_base") {
+    shader_override = graph.add_scalar<int64_t>(kShaderOverrideForceBase);
   } else if (impl_selector == "non_gqa") {
-    shader_override = graph.add_scalar<int64_t>(0);
+    shader_override = graph.add_scalar<int64_t>(kShaderOverrideForceNonGqa);
   }
 
   test_sdpa_impl(

backends/vulkan/test/custom_ops/test_sdpa.cpp

@@ -309,8 +309,8 @@ static std::vector<TestCase> generate_sdpa_test_cases() {
   // attn_weights S/context alignment (a decode-shaped buffer allocation has no
   // headroom for the shaders' align_up_4 stride unless padded — see sdpa_impl).
   {
-    // Cover both D=64 (D4=16) and D=128 (D4=32) so the GQA AV path is
-    // validated across head_dim.
+    // Cover D=64 (D4=16) and D=128 (D4=32) with the vendor-default GQA and the
+    // per-query-head shaders, across texture + buffer.
     const std::vector<SDPAConfig> decs = {
         {64, 8, 2, 1, 32, "accu", "decode"},
         {128, 8, 2, 1, 32, "accu_d128", "decode"},
@@ -323,6 +323,23 @@ static std::vector<TestCase> generate_sdpa_test_cases() {
         }
       }
     }
+
+    // Force the head_dim output-tiled GQA variant (Adreno-only in production)
+    // so its wg x-collapse and the partial_n_tile tail get deterministic
+    // coverage on any device: D=64/128 give even D4 (fast path); D=4 gives D4=1
+    // (odd), exercising the partial-tile checked load.
+    const std::vector<SDPAConfig> tile2_decs = {
+        {64, 8, 2, 1, 32, "accu_tile2", "decode"},
+        {128, 8, 2, 1, 32, "accu_tile2_d128", "decode"},
+        {4, 8, 2, 1, 32, "accu_tile2_d4", "decode"},
+    };
+    for (const auto& dec : tile2_decs) {
+      for (const auto& storage : {utils::kTexture3D, utils::kBuffer}) {
+        test_cases.push_back(
+            create_sdpa_test_case(dec, vkapi::kFloat, storage, "gqa_tile2"));
+      }
+    }
+
     SDPAConfig pre{64, 8, 2, 16, 16, "accu", "prefill"};
     test_cases.push_back(create_sdpa_test_case(
         pre, vkapi::kFloat, utils::kTexture3D, "default"));