diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml index 6cf919910d5..546eb9da7c0 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml +++ b/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 diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 83891d5ed74..c4259d4c530 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/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,7 +201,8 @@ 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(shader_override) != 0; + const bool force_gqa = graph->extract_scalar(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. @@ -208,12 +210,29 @@ bool resolve_use_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(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& 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(-2, v); - return {N4, M4, static_cast(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(num_kv_heads * d.B)}; } return {N4, M4, static_cast(d.H * d.B)}; diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.h b/backends/vulkan/runtime/graph/ops/impl/SDPA.h index b40841126b8..fa51ebef111 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.h +++ 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, diff --git a/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp b/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp index 38e674faeee..4a59210cf9c 100644 --- a/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp +++ b/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& 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(-3, q); const int64_t context_len = graph.size_at(-3, k_cache); @@ -147,12 +151,17 @@ void test_sdpa(ComputeGraph& graph, const std::vector& 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(1); + shader_override = graph.add_scalar(kShaderOverrideForceGqa); + } else if (impl_selector == "gqa_tile2") { + shader_override = graph.add_scalar(kShaderOverrideForceTile2); + } else if (impl_selector == "gqa_base") { + shader_override = graph.add_scalar(kShaderOverrideForceBase); } else if (impl_selector == "non_gqa") { - shader_override = graph.add_scalar(0); + shader_override = graph.add_scalar(kShaderOverrideForceNonGqa); } test_sdpa_impl( diff --git a/backends/vulkan/test/custom_ops/test_sdpa.cpp b/backends/vulkan/test/custom_ops/test_sdpa.cpp index 414363e0413..55d2173e3cd 100644 --- a/backends/vulkan/test/custom_ops/test_sdpa.cpp +++ b/backends/vulkan/test/custom_ops/test_sdpa.cpp @@ -309,8 +309,8 @@ static std::vector 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 decs = { {64, 8, 2, 1, 32, "accu", "decode"}, {128, 8, 2, 1, 32, "accu_d128", "decode"}, @@ -323,6 +323,23 @@ static std::vector 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 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"));