From 145e0c90ddfa0d40249d0231bb7f7f6fc01a67d7 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Tue, 21 Jul 2026 16:55:52 -0700 Subject: [PATCH 1/2] [ET-VK][sdpa] Reuse shared V cache across GQA query heads in AV coop-GEMV Pull Request resolved: https://github.com/pytorch/executorch/pull/21063 The LLM decode AV coop-GEMV reloads the shared V cache once per query head. In grouped-query attention Hq = G * Hkv query heads share each KV head (Llama G=4, Phi G=3, Qwen G=2), and out[q_h, d] = sum_c attn[c, q_h] * V[c, kv_h, d] reads the SAME V texel for every query head in a group. The per-query-head coop shader gives each of the Hq heads its own workgroup, so V -- the dominant traffic (head_dim-wide per context texel, vs a scalar attn weight) -- is read G times. This adds a GQA-reuse AV variant that assigns ONE workgroup per (d4, kv_h): it loads each V texel once and reuses it across all G query heads in the group, producing G output texels. For this bandwidth-bound kernel that cuts V-cache traffic ~Gx. Implementation: - The variant is a codegen flag (`GQA`) on the existing `sdpa_compute_out_coop.glsl` template, not a separate file: one shared header plus two `#ifdef GQA` `main()`s (per-head and GQA-reuse), so the shared setup lives in one place while each algorithm reads end-to-end. It emits the shader `sdpa_compute_out_gqa_coop`. - Reduction reuses the per-head coop shader's shared-memory tree reduction (no subgroup arithmetic), so the variant runs on any Vulkan device -- Adreno and Mali alike -- with no capability gate. - Each thread holds G output accumulators; the array is sized to a compile-time `MAX_GROUP_SIZE` = 8 and the group loop is bounded by the `group_size` = Hq/Hkv spec constant, so the driver fully unrolls it at pipeline creation. - Dispatch (`pick_sdpa_av_shader` + global-wg picker + spec-const wiring in `add_sdpa_compute_out_node`): the GQA variant is selected on the LLM decode coop path when Hq > Hkv, evenly divisible, and G <= 8 (`use_gqa_av_coop`); it sets `group_size` and changes the global workgroup z-dim from Hq to Hkv. Everything else -- MHA (Hq == Hkv), groups exceeding the cap (G > 8, e.g. MQA with Hq > 8), and non-divisible shapes -- falls back to the unchanged per-head `sdpa_compute_out_coop`. (Low-ratio MQA -- Hkv == 1 with Hq <= 8 -- is eligible and takes the GQA path.) - A test-only `gqa_override` knob is threaded through `add_sdpa_compute_out_node` (declared in the new `SDPA.h`): -1 auto-select, 0 force per-head, 1 force GQA, so a benchmark can exercise both AV shaders on the same shape; forcing GQA is VK_CHECK'd against shape eligibility. ghstack-source-id: 405400503 @exported-using-ghexport Differential Revision: [D112906311](https://our.internmc.facebook.com/intern/diff/D112906311/) --- .../graph/ops/glsl/sdpa_compute_out_coop.glsl | 185 ++++++++++++++++++ .../graph/ops/glsl/sdpa_compute_out_coop.yaml | 3 + .../vulkan/runtime/graph/ops/impl/SDPA.cpp | 157 ++++++++++++--- backends/vulkan/runtime/graph/ops/impl/SDPA.h | 68 +++++++ .../vulkan/test/custom_ops/impl/TestSDPA.cpp | 138 ++++++++++--- backends/vulkan/test/custom_ops/test_sdpa.cpp | 49 +++-- backends/vulkan/test/op_tests/sdpa_test.cpp | 13 ++ 7 files changed, 547 insertions(+), 66 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/impl/SDPA.h diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl index cd2c689ebc8..f6f00c9bfe5 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl @@ -32,6 +32,16 @@ $if V_CACHE_STORAGE == "buffer": #define NUM_WORKERS_PER_OUT 64 +$if GQA: + #define GQA + // Maximum grouped-query-attention group size (Hq / Hkv) supported by the GQA + // variant. The actual group size G is passed as a specialization constant and + // must satisfy G <= MAX_GROUP_SIZE. The accumulator array is sized to + // MAX_GROUP_SIZE (a compile-time constant) and the group loop is bounded by + // the spec-const G so the driver can fully unroll it at pipeline creation + // time. Group sizes seen in practice: Llama G=4, Phi G=3, Qwen G=2. + #define MAX_GROUP_SIZE 8 + ${define_required_extensions(IO_STORAGE, DTYPE)} layout(std430) buffer; @@ -49,6 +59,8 @@ ${layout_declare_ubo(B, "int", "input_pos")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} +$if GQA: + ${layout_declare_spec_const(C, "int", "group_size", "1")} #include "sdpa_fp_attn_weight_tile_load.glslh" #include "sdpa_fp_v_cache_tile_load.glslh" @@ -57,6 +69,177 @@ ${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} shared FPOutTile partial_sums[NUM_WORKERS_PER_OUT]; +#ifdef GQA + +/* + * GQA-reuse variant of the AV coop GEMV. + * + * Grouped-query attention shares one KV head across G = Hq / Hkv query heads. + * The AV computation out[q_h, d] = sum_c attn[c, q_h] * V[c, kv_h(q_h), d] + * reads the SAME V cache for every query head in a group. In the per-query-head + * coop variant each of the Hq heads gets its own workgroup and independently + * re-loads the shared V cache, so V (the dominant traffic — head_dim wide per + * context texel) is read G times. This variant assigns ONE workgroup per + * (d4, kv_h): it loads each V texel once and reuses it across all G query heads + * in the group, producing G output texels. This cuts V-cache traffic ~Gx for a + * bandwidth-bound kernel. Dispatch sets the global wg z-dim to Hkv (not Hq). + */ + +void main() { + const int worker_id = int(gl_LocalInvocationID.y); + + const int tile_idx_x = int(gl_GlobalInvocationID.x); + // idx along the K/V head dim + const int kv_h = int(gl_GlobalInvocationID.z); + + // idx along the output head_dim dim + const int d = tile_idx_x * TILE_N; + const int d4 = div_4(d); + + // idx along the output seq_len dim. Note that for this shader seq_len will be + // 1. + const int s = 0; + + // texel size of head_dim + const int D4 = div_up_4(q_sizes.x); + // number of Q heads + const int Q_H = q_sizes.y; + // sequence length + const int S = q_sizes.z; + const int S_aligned = align_up_4(S); + + // number of K/V heads + const int KV_H = v_sizes.y; + // Max context length + const int C = v_sizes.z; + const int C4 = div_up_4(C); + + const int G = group_size; + // First query head in this group. + const int q_h_base = kv_h * G; + + // current context length + const int context_len = input_pos + S; + const int context_texel_len = div_up_4(context_len); + + // bounds check + if (d4 >= D4 || s >= S || kv_h >= KV_H) { + return; + } + + // With head_dim output-tiling (TILE_N4 > 1) each workgroup owns TILE_N4 + // head_dim texels. Whether this workgroup owns a partial tile (only possible + // on the final tile when D4 % TILE_N4 != 0). Uniform across the workgroup, so + // the branch selecting the checked vs unchecked V load below is a uniform + // branch. For the common even-D4 case (D=64 -> D4=16, D=128 -> D4=32, and + // always for TILE_N4 == 1) this is false and the fast unchecked path is taken. + const bool partial_n_tile = (d4 + TILE_N4) > D4; + + FPOutTile out_tile[MAX_GROUP_SIZE]; + [[unroll]] for (int g = 0; g < MAX_GROUP_SIZE; ++g) { + if (g >= G) { + break; + } + initialize(out_tile[g]); + } + + FPInputTile attn_weight_tile; + FPWeightTile w_tile; + + const int context_len_aligned_down = context_len - mod_4(context_len); + const int C4_limit = div_up_4(context_len_aligned_down); + + // Main loop: each thread strides over context texels. The single V load per + // texel is reused across all G query heads in the group. + for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_OUT) { + const int c = mul_4(c4); + + // The TILE_N4 V texels are loaded once per context texel and reused across + // all G query heads in the group. Bounds-check the head_dim index only on a + // partial final tile (uniform branch, see partial_n_tile). + if (partial_n_tile) { + load_v_cache_tile_with_checks( + w_tile, d4, c, kv_h, D4, context_len, C, KV_H); + } else { + load_v_cache_tile_no_checks( + w_tile, d4, c, kv_h, D4, context_len, C, KV_H); + } + + [[unroll]] for (int g = 0; g < MAX_GROUP_SIZE; ++g) { + if (g >= G) { + break; + } + load_attn_weight_tile_no_checks( + attn_weight_tile, + c4, + s, + q_h_base + g, + context_texel_len, + S_aligned, + Q_H); + fp_accumulate_with_fp_weight(out_tile[g], attn_weight_tile, w_tile); + } + } + + // first worker in the work group will handle final texel, which may contain + // padding elements. + if (worker_id == 0) { + for (int c4 = C4_limit; c4 < context_texel_len; c4++) { + const int c = mul_4(c4); + + load_v_cache_tile_with_checks( + w_tile, d4, c, kv_h, D4, context_len, C, KV_H); + + [[unroll]] for (int g = 0; g < MAX_GROUP_SIZE; ++g) { + if (g >= G) { + break; + } + load_attn_weight_tile_with_checks( + attn_weight_tile, + c4, + s, + q_h_base + g, + context_texel_len, + S_aligned, + Q_H); + fp_accumulate_with_fp_weight(out_tile[g], attn_weight_tile, w_tile); + } + } + } + + // Combine the per-worker partial sums for each group with a shared-memory tree + // reduction. partial_sums is reused across groups; no trailing barrier is + // needed because each worker only writes its own slot and the next group's + // leading barrier orders those writes after this group's reduction reads + // (worker 0's store reads only slot 0, which only worker 0 writes). + [[unroll]] for (int g = 0; g < MAX_GROUP_SIZE; ++g) { + if (g >= G) { + break; + } + partial_sums[worker_id] = out_tile[g]; + + memoryBarrierShared(); + barrier(); + + for (int i = NUM_WORKERS_PER_OUT / 2; i > 0; i /= 2) { + if (worker_id < i) { + accumulate_out_tile_with_out_tile( + partial_sums[worker_id], partial_sums[worker_id + i]); + } + memoryBarrierShared(); + barrier(); + } + + if (worker_id == 0) { + out_tile[g] = partial_sums[0]; + store_sdpa_out_tile_with_checks( + out_tile[g], d4, s, q_h_base + g, D4, S, Q_H); + } + } +} + +#else + /* * See the tiled variant of this shader for the implemented behavior. This * shader is implements an optimization for cases where sequence length is 1; in @@ -197,3 +380,5 @@ void main() { Q_H); } } + +#endif // GQA 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 33ec2f8b322..6cf919910d5 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 @@ -11,6 +11,7 @@ sdpa_compute_out_coop: V_CACHE_STORAGE: texture3d TILE_K4: 1 TILE_N4: 1 + GQA: False generate_variant_forall: combination: parameter_names: [IO_STORAGE, V_CACHE_STORAGE] @@ -23,3 +24,5 @@ sdpa_compute_out_coop: - VALUE: half shader_variants: - NAME: sdpa_compute_out_coop + - NAME: sdpa_compute_out_gqa_coop + GQA: True diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 3efb834725d..83891d5ed74 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -27,19 +28,6 @@ namespace vkcompute { namespace { -// -// SDPA mode: distinguishes the two dispatch families sharing this file. -// LLM — Llama-style KV-cache SDPA. Q layout [B=1, S, H, D] (DHSB). -// Separate k_cache/v_cache inputs + input_pos_symint for dynamic -// context_len. attn_weights are padded to multiples of 4 in the -// S/context_len dims and carry the input dtype. A coop (GEMV) -// shader variant is selected for single-token decode. -// FUSED — General SDPA fused op. Q layout [B, H, S, D] (DSHB). No cache, -// optional additive attn_mask, optional scale arg. attn_weights -// are unpadded and always fp32. Tiled shader variant only. -// -enum class SDPAMode { LLM, FUSED }; - // // Common dimension helper: folds the axis-swap for LLM vs fused Q layouts. // `input_pos_symint` is used only for LLM (context_len = S + input_pos); @@ -178,6 +166,54 @@ static inline SDPAMode mode_of(const std::vector& resize_args) { return static_cast(resize_args.at(3)); } +// Upper bound on the GQA group size the coop-GQA shader supports; must match +// MAX_GROUP_SIZE in the GQA branch of sdpa_compute_out_coop.glsl. +constexpr int64_t kMaxGqaGroupSize = 8; + +// Whether the LLM decode AV path should use the GQA-reuse coop shader: one +// workgroup computes all G = Hq / Hkv query heads sharing a KV head, loading +// each V texel once. Requires GQA (Hq > Hkv, evenly divisible) and a group size +// within the shader's compile-time bound. +bool use_gqa_av_coop( + ComputeGraph* graph, + const int64_t num_q_heads, + const int64_t num_kv_heads) { + (void)graph; + if (num_kv_heads <= 0 || num_q_heads <= num_kv_heads) { + return false; + } + if (num_q_heads % num_kv_heads != 0) { + return false; + } + const int64_t group_size = num_q_heads / num_kv_heads; + return group_size <= kMaxGqaGroupSize; +} + +// 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. +bool resolve_use_gqa( + ComputeGraph* graph, + const ValueRef shader_override, + const int64_t num_q_heads, + const int64_t num_kv_heads) { + 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; + // 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"); + } + return force_gqa; +} + vkapi::ShaderInfo pick_sdpa_qk_shader( ComputeGraph* graph, const std::vector& args, @@ -297,8 +333,26 @@ vkapi::ShaderInfo pick_sdpa_av_shader( const ValueRef q_projected = resize_args.at(0); const bool is_gemv = is_single_token(graph, q_projected); + // Test-only knob (see SDPA.h): -1 auto, 0 forces per-query-head, 1 forces + // the GQA-reuse shader. + const ValueRef shader_override = resize_args.at(4); + std::string shader_name = "sdpa_compute_out"; - shader_name += is_gemv ? "_coop" : "_tiled"; + if (is_gemv) { + const int64_t num_q_heads = graph->size_at(-2, q_projected); + const int64_t num_kv_heads = graph->size_at(-2, v_cache); + if (resolve_use_gqa(graph, shader_override, num_q_heads, num_kv_heads)) { + // Grouped-query attention: one workgroup computes all G = Hq / Hkv + // query heads that share a KV head, loading each V texel once and + // reusing it across the group. Cuts the dominant V-cache traffic ~Gx + // for this bandwidth-bound kernel. + shader_name += "_gqa_coop"; + } else { + shader_name += "_coop"; + } + } else { + shader_name += "_tiled"; + } add_storage_type_suffix(shader_name, graph->storage_type_of(out)); add_storage_type_suffix(shader_name, graph->storage_type_of(v_cache)); add_dtype_suffix(shader_name, graph->dtype_of(out)); @@ -319,7 +373,6 @@ utils::uvec3 pick_sdpa_av_global_wg_size( const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { - (void)shader; const SDPAMode mode = mode_of(resize_args); const ValueRef q = resize_args.at(0); const ValueRef k = resize_args.at(1); @@ -328,6 +381,17 @@ utils::uvec3 pick_sdpa_av_global_wg_size( const uint32_t N4 = utils::div_up_4(static_cast(d.D)); const uint32_t M4 = utils::div_up_4(static_cast(d.S)); + + // The GQA-reuse AV coop shader assigns one workgroup per (d4, kv_h) and emits + // all G query heads in the group, so the z-dim is the number of KV heads + // rather than the number of Q heads. d.H is the Q-head count; recover the KV + // count from the V cache (args read group: [attn_weights_softmax, v]). + 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)}; + } + return {N4, M4, static_cast(d.H * d.B)}; } @@ -498,7 +562,8 @@ void add_sdpa_compute_out_node( const ValueRef k, const ValueRef input_pos_symint, const ValueRef out, - const SDPAMode mode) { + const SDPAMode mode, + const ValueRef shader_override) { vkapi::ParamsBindList param_ubos; if (mode == SDPAMode::LLM) { param_ubos = { @@ -511,6 +576,24 @@ void add_sdpa_compute_out_node( const ValueRef mode_ref = static_cast(mode); + // The GQA-reuse AV coop shader needs its group size G = Hq / Hkv as a + // specialization constant. The shader declares spec consts [inv_scale (unused + // in AV), group_size]; pass both positionally so group_size lands at the + // right slot. Non-GQA AV shaders declare only inv_scale (or none), so the + // trailing group_size entry is simply ignored there. The GQA gate is fixed at + // build time (Hq/Hkv/capability don't change across decode steps), matching + // pick_sdpa_av_shader. + vkapi::SpecVarList spec_vars = {}; + if (mode == SDPAMode::LLM) { + const int64_t num_q_heads = graph.size_at(-2, q); + const int64_t num_kv_heads = graph.size_at(-2, v); + if (resolve_use_gqa(&graph, shader_override, num_q_heads, num_kv_heads)) { + const int32_t group_size = + utils::safe_downcast(num_q_heads / num_kv_heads); + spec_vars = {1.0f, group_size}; + } + } + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_sdpa_av_shader, @@ -523,9 +606,9 @@ void add_sdpa_compute_out_node( // Push Constants {}, // Specialization Constants - {}, - // Resize Args: [q, k, input_pos_symint_or_dummy, mode] - {q, k, input_pos_symint, mode_ref}, + spec_vars, + // Resize Args: [q, k, input_pos_symint_or_dummy, mode, shader_override] + {q, k, input_pos_symint, mode_ref, shader_override}, // Resizing Logic resize_sdpa_out_node)); } @@ -599,17 +682,31 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { const utils::StorageType attn_weights_storage = graph.storage_type_of(q_projected); - // If using buffer storage for attn weights, we need to ensure that the buffer - // numel limit is not exceeded. If needed, manually adjust max_seq_len based - // on the buffer numel limit. + // The attn_weights S and context dims are padded up to a multiple of 4 to + // match resize_sdpa_attn_weights_node(), which virtual_resizes to + // align_up_4(seq_len)/align_up_4(context_len). The alignment originates from + // an Adreno 750 (Samsung S24) correctness workaround (D86226134 / PR #15578). + // With buffer-backed attn_weights the SDPA shaders index each head at a + // per-head stride of align_up_4(S) * div_up_4(context) texels, so the + // allocation must reserve that padded extent — otherwise heads past the first + // few read/write out of bounds. The texture path is immune: it addresses each + // head as an image layer rather than via a linear stride. + const int64_t padded_context_len = utils::align_up_4(max_context_len); + + // If using buffer storage for attn weights, ensure the buffer numel limit is + // not exceeded by the PADDED allocation. Checking the raw (unpadded) sizes is + // insufficient: padding can push the actual allocation past the limit even + // when the raw size fits — e.g. max_seq_len=2047 pads to 2048, so + // 32*2048*2048 lands exactly on maxStorageBufferRange/4 and the shaders index + // out of bounds, corrupting attention output. if (attn_weights_storage == utils::kBuffer) { const int64_t max_buffer_numel = graph.max_buffer_numel(); - if (num_q_heads * max_seq_len * max_context_len >= max_buffer_numel) { - // Compute the maximum possible value for max_seq_len that will hit - // the buffer numel limit. - max_seq_len = max_buffer_numel / (num_q_heads * max_context_len); - // Adjust down to the nearest multiple of 4 to make sure the limit is - // not hit. + if (num_q_heads * utils::align_up_4(max_seq_len) * padded_context_len >= + max_buffer_numel) { + // Largest max_seq_len whose padded allocation stays under the limit. + max_seq_len = max_buffer_numel / (num_q_heads * padded_context_len); + // Round down to a multiple of 4 (so align_up_4 is a no-op below) and keep + // the allocation strictly below the limit. if (max_seq_len % 4 != 0) { max_seq_len = (max_seq_len / 4) * 4; } else { @@ -621,8 +718,8 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { std::vector attn_weight_full_sizes = { 1, // batch num_q_heads, - max_seq_len, - max_context_len}; + utils::align_up_4(max_seq_len), + padded_context_len}; TmpTensor attn_weights( &graph, diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.h b/backends/vulkan/runtime/graph/ops/impl/SDPA.h new file mode 100644 index 00000000000..b40841126b8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.h @@ -0,0 +1,68 @@ +/* + * 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 + +namespace vkcompute { + +// +// SDPA mode: distinguishes the two dispatch families sharing SDPA.cpp. +// LLM — Llama-style KV-cache SDPA. Q layout [B=1, S, H, D] (DHSB). +// Separate k_cache/v_cache inputs + input_pos_symint for dynamic +// context_len. attn_weights are padded to multiples of 4 in the +// S/context_len dims and carry the input dtype. A coop (GEMV) +// shader variant is selected for single-token decode. +// FUSED — General SDPA fused op. Q layout [B, H, S, D] (DSHB). No cache, +// optional additive attn_mask, optional scale arg. attn_weights +// are unpadded and always fp32. Tiled shader variant only. +// +enum class SDPAMode { LLM, FUSED }; + +void add_sdpa_compute_attn_weights_node( + ComputeGraph& graph, + const ValueRef q, + const ValueRef k, + const ValueRef input_pos_symint, + const ValueRef attn_mask, + const float scale_val, + const ValueRef attn_weights, + const SDPAMode mode); + +void add_sdpa_attn_weights_softmax_node( + ComputeGraph& graph, + const ValueRef attn_weights, + const ValueRef q, + const ValueRef k, + const ValueRef input_pos_symint, + const ValueRef attn_weights_softmax, + const SDPAMode mode); + +// 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. +void add_sdpa_compute_out_node( + ComputeGraph& graph, + const ValueRef attn_weights_softmax, + const ValueRef v, + const ValueRef q, + const ValueRef k, + const ValueRef input_pos_symint, + const ValueRef out, + const SDPAMode mode, + const ValueRef shader_override = kDummyValueRef); + +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp b/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp index dc908de55e6..38e674faeee 100644 --- a/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp @@ -9,9 +9,104 @@ #include #include +#include + +#include namespace vkcompute { +// Bare-minimum mirror of the LLM decode/prefill path in sdpa_impl() (see +// SDPA.cpp), stripped of the input-validation VK_CHECK_CONDs. Building the +// three SDPA nodes directly lets the test forward a shader_override knob to the +// AV node, which the production op boundary (llama.custom_sdpa) can't carry. +static void test_sdpa_impl( + ComputeGraph& graph, + const ValueRef q_projected, + const ValueRef k_cache, + const ValueRef v_cache, + const ValueRef input_pos_symint, + const ValueRef out, + const ValueRef shader_override) { + const int64_t num_q_heads = graph.size_at(-2, q_projected); + int64_t max_seq_len = graph.size_at(-3, q_projected); + const int64_t max_context_len = graph.size_at(-3, k_cache); + + const utils::StorageType attn_weights_storage = + graph.storage_type_of(q_projected); + + // S and context dims are padded to a multiple of 4 to match sdpa_impl() (see + // the allocation comment there): the buffer SDPA shaders index each head at + // an align_up_4(S)/align_up_4(context) stride, so the allocation must reserve + // the padded extent or later heads land out of bounds on the buffer path. + const int64_t padded_context_len = utils::align_up_4(max_context_len); + + // Clamp attn-weight seq_len so the PADDED buffer stays within the numel limit + // (mirrors sdpa_impl; a no-op for the texture path the harness uses). The + // check must use the padded sizes actually allocated, not the raw ones. + if (attn_weights_storage == utils::kBuffer) { + const int64_t max_buffer_numel = graph.max_buffer_numel(); + if (num_q_heads * utils::align_up_4(max_seq_len) * padded_context_len >= + max_buffer_numel) { + max_seq_len = max_buffer_numel / (num_q_heads * padded_context_len); + if (max_seq_len % 4 != 0) { + max_seq_len = (max_seq_len / 4) * 4; + } else { + max_seq_len -= 4; + } + } + } + + const std::vector attn_weight_full_sizes = { + 1, num_q_heads, utils::align_up_4(max_seq_len), padded_context_len}; + + TmpTensor attn_weights( + &graph, + attn_weight_full_sizes, + graph.dtype_of(q_projected), + attn_weights_storage, + utils::kWidthPacked); + + TmpTensor attn_weights_softmax( + &graph, + attn_weight_full_sizes, + graph.dtype_of(q_projected), + attn_weights_storage, + utils::kWidthPacked); + + const int32_t head_dim_size = graph.size_at(-1, q_projected); + const float scale_val = 1.0f / std::sqrt(static_cast(head_dim_size)); + + add_sdpa_compute_attn_weights_node( + graph, + q_projected, + k_cache, + input_pos_symint, + /*attn_mask=*/kDummyValueRef, + scale_val, + attn_weights, + SDPAMode::LLM); + + add_sdpa_attn_weights_softmax_node( + graph, + attn_weights, + q_projected, + k_cache, + input_pos_symint, + attn_weights_softmax, + SDPAMode::LLM); + + add_sdpa_compute_out_node( + graph, + attn_weights_softmax, + v_cache, + q_projected, + /*k=*/kDummyValueRef, + input_pos_symint, + out, + SDPAMode::LLM, + shader_override); +} + // Test wrapper for the LLM KV-cache SDPA op (llama.custom_sdpa.default). // // The production op reads the dynamic context length from an `input_pos` @@ -19,14 +114,15 @@ namespace vkcompute { // context_len = S + input_pos, and context_len is exactly the KV-cache's dim // -3. Since q is [B=1, S, H, D] and the caches are [B=1, context_len, H_kv, D], // input_pos is fully determined by the tensor shapes. Deriving it here from -// those shapes keeps the single source of truth (the cache size) authoritative; -// exposing input_pos as its own ValueSpec would add a redundant degree of -// freedom that a test could set inconsistent with the cache shape. +// those shapes keeps the single source of truth (the cache size) authoritative. // -// Decode vs prefill is selected automatically inside the op via -// is_single_token() (S == 1 -> coop/GEMV shaders, S > 1 -> tiled shaders), so -// impl_selector is accepted for interface uniformity but only "default" is -// meaningful here. +// 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 +// impl_selector has no effect on prefill (tiled). // // Args: q, k_cache, v_cache, impl_selector, out void test_sdpa(ComputeGraph& graph, const std::vector& args) { @@ -39,8 +135,9 @@ 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", - "test_sdpa only supports the 'default' impl_selector"); + impl_selector == "default" || impl_selector == "gqa" || + impl_selector == "non_gqa", + "test_sdpa: impl_selector must be one of {default, gqa, non_gqa}"); const int64_t seq_len = graph.size_at(-3, q); const int64_t context_len = graph.size_at(-3, k_cache); @@ -50,19 +147,16 @@ void test_sdpa(ComputeGraph& graph, const std::vector& args) { const ValueRef input_pos_symint = graph.add_symint(input_pos_val); - VK_GET_OP_FN("llama.custom_sdpa.default") - (graph, - { - q, - k_cache, - v_cache, - input_pos_symint, - kDummyValueRef, // attn_mask - kDummyValueRef, // dropout_p - kDummyValueRef, // is_causal (implementation assumes causal) - kDummyValueRef, // scale (implementation computes 1/sqrt(head_dim)) - out, - }); + // shader_override (see SDPA.h): -1 auto, 0 force non-GQA, 1 force GQA. + ValueRef shader_override = kDummyValueRef; + if (impl_selector == "gqa") { + shader_override = graph.add_scalar(1); + } else if (impl_selector == "non_gqa") { + shader_override = graph.add_scalar(0); + } + + test_sdpa_impl( + graph, q, k_cache, v_cache, input_pos_symint, out, shader_override); } REGISTER_OPERATORS { diff --git a/backends/vulkan/test/custom_ops/test_sdpa.cpp b/backends/vulkan/test/custom_ops/test_sdpa.cpp index 15040c781db..414363e0413 100644 --- a/backends/vulkan/test/custom_ops/test_sdpa.cpp +++ b/backends/vulkan/test/custom_ops/test_sdpa.cpp @@ -54,7 +54,6 @@ struct SDPAConfig { int64_t context_len; // total KV length (kv_len) std::string model; // label only std::string regime; // "decode" / "prefill", label only - std::string impl_selector = "default"; }; static std::vector as_float_data(const ValueSpec& spec) { @@ -75,7 +74,8 @@ static std::vector as_float_data(const ValueSpec& spec) { static TestCase create_sdpa_test_case( const SDPAConfig& config, vkapi::ScalarType dtype, - utils::StorageType storage_type) { + utils::StorageType storage_type, + const std::string& impl) { TestCase test_case; const bool is_perf = config.context_len > kRefContextLenLimit; @@ -89,7 +89,8 @@ static TestCase create_sdpa_test_case( std::to_string(config.seq_len) + " C" + std::to_string(config.context_len); - const std::string suffix = "[" + config.model + " " + config.regime + "]"; + const std::string suffix = + "[" + config.model + " " + config.regime + " " + impl + "]"; test_case.set_name(make_test_label( prefix, dtype_str, dtype_str, shape, storage_str, suffix)); @@ -117,7 +118,7 @@ static TestCase create_sdpa_test_case( utils::kWidthPacked, DataGenType::RANDOM); - ValueSpec impl_selector = ValueSpec::make_string(config.impl_selector); + ValueSpec impl_selector = ValueSpec::make_string(impl); // out: [1, S, n_heads, head_dim] ValueSpec output( @@ -263,6 +264,11 @@ static std::vector generate_sdpa_test_cases() { const auto dtype = vkapi::kHalf; const auto storage = utils::kTexture3D; + // Decode (S==1) picks a coop AV shader; exercise both the GQA-reuse variant + // and the per-query-head variant for every decode case. Prefill (tiled) is + // unaffected by the selector, so it runs a single case. + const std::vector decode_impls = {"gqa", "non_gqa"}; + for (const auto& m : models) { for (int64_t c : decode_context_lens) { SDPAConfig cfg; @@ -273,7 +279,9 @@ static std::vector generate_sdpa_test_cases() { cfg.context_len = c; cfg.model = m.name; cfg.regime = "decode"; - test_cases.push_back(create_sdpa_test_case(cfg, dtype, storage)); + for (const auto& impl : decode_impls) { + test_cases.push_back(create_sdpa_test_case(cfg, dtype, storage, impl)); + } } if (decode_only_mode()) { continue; @@ -287,7 +295,8 @@ static std::vector generate_sdpa_test_cases() { cfg.context_len = s; cfg.model = m.name; cfg.regime = "prefill"; - test_cases.push_back(create_sdpa_test_case(cfg, dtype, storage)); + test_cases.push_back( + create_sdpa_test_case(cfg, dtype, storage, "default")); } } @@ -295,16 +304,28 @@ static std::vector generate_sdpa_test_cases() { return test_cases; } - // Small ACCU correctness cases (fp32, texture), decode + prefill. Texture is - // the production LLM path; the buffer decode path has a known GQA-head - // discrepancy that is out of scope for this perf harness. + // Small ACCU correctness cases (fp32), decode + prefill. Texture is the + // production LLM path; buffer is also validated for decode to guard the + // 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). { - SDPAConfig dec{64, 8, 2, 1, 32, "accu", "decode"}; - test_cases.push_back( - create_sdpa_test_case(dec, vkapi::kFloat, utils::kTexture3D)); + // Cover both D=64 (D4=16) and D=128 (D4=32) so the GQA AV path is + // validated across head_dim. + const std::vector decs = { + {64, 8, 2, 1, 32, "accu", "decode"}, + {128, 8, 2, 1, 32, "accu_d128", "decode"}, + }; + for (const auto& dec : decs) { + for (const auto& storage : {utils::kTexture3D, utils::kBuffer}) { + for (const auto& impl : decode_impls) { + test_cases.push_back( + create_sdpa_test_case(dec, vkapi::kFloat, storage, impl)); + } + } + } SDPAConfig pre{64, 8, 2, 16, 16, "accu", "prefill"}; - test_cases.push_back( - create_sdpa_test_case(pre, vkapi::kFloat, utils::kTexture3D)); + test_cases.push_back(create_sdpa_test_case( + pre, vkapi::kFloat, utils::kTexture3D, "default")); } return test_cases; diff --git a/backends/vulkan/test/op_tests/sdpa_test.cpp b/backends/vulkan/test/op_tests/sdpa_test.cpp index fd2afc7408e..c1c252165eb 100644 --- a/backends/vulkan/test/op_tests/sdpa_test.cpp +++ b/backends/vulkan/test/op_tests/sdpa_test.cpp @@ -879,3 +879,16 @@ TEST(VulkanSDPATest, test_sdpa_op_llama3_params_dynamic) { test_vulkan_sdpa( 0, {111, 1, 1, 1, 57, 1, 1}, head_dim, num_heads, num_kv_heads, 1); } + +// GQA group size G = Hq / Hkv = 4 (Llama-style). The other decode tests cover +// G=2 (small_params: 8/4) and G=3 (small_params_dynamic: 6/2, llama3: 24/8); +// this exercises the G=4 path of the AV coop-GQA shader, which reuses each V +// texel across all query heads in a group. +TEST(VulkanSDPATest, test_sdpa_op_gqa_group4) { + const int head_dim = 64; + const int num_heads = 8; + const int num_kv_heads = 2; + + test_vulkan_sdpa( + 0, {5, 1, 1, 1, 1, 3, 1}, head_dim, num_heads, num_kv_heads, 1); +} From 26c762b957e80380bcee50cdfcbe86e72946ddab Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Tue, 21 Jul 2026 16:55:52 -0700 Subject: [PATCH 2/2] [ET-VK][sdpa] Vendor-adaptive head_dim output-tiling (TILE_N4=2) in GQA AV coop-GEMV MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21064 Builds on the GQA-reuse AV coop-GEMV (parent commit). The base `sdpa_compute_out_gqa_coop` shader assigns one head_dim texel (TILE_N4=1) per workgroup x-slot. This adds a head_dim output-tiled variant, `sdpa_compute_out_gqa_coop_tile2` (TILE_N4=2): each workgroup owns 2 head_dim texels and emits G x TILE_N4 outputs, so the per-context-texel attn-weight loads and the shared-memory tree reduction are amortized over twice as many outputs. The AV coop GEMV GLSL is already a TILE_N4-parameterized template (the `partial_n_tile` uniform branch handles the D4 % TILE_N4 != 0 tail), so this change is purely a new codegen variant plus dispatch wiring — no shader-body change. Selection is vendor-adaptive: `pick_sdpa_av_shader` appends the `_tile2` suffix only on Adreno (`graph->device_is_adreno()`). Tiling is a consistent win on Adreno (AV ~1.14-1.63x over the base GQA variant) but a regression on Mali at common decode contexts (~0.67-0.86x, interleaved median-of-N), so Mali and other vendors keep the base `sdpa_compute_out_gqa_coop`. `pick_sdpa_av_global_wg_size` keys off the same `_tile2` suffix: the tiled variant's x-dim collapses from D4 to div_up(D4, 2) since each workgroup now covers 2 head_dim texels. Because the variant is Adreno-only in production, the test-only `gqa_override` knob (see SDPA.h) is extended so tests can pin the variant on any device: `kGqaOverrideForceTile2` / `kGqaOverrideForceBase` force the tiled / base variant regardless of vendor (via `resolve_use_tile2`), giving the tiled shader deterministic coverage on Mali / SwiftShader. `resolve_use_gqa` is unchanged (any non-`ForceNonGqa` value still forces the GQA family, VK_CHECK'd for eligibility). ghstack-source-id: 405400514 @exported-using-ghexport Differential Revision: [D112906313](https://our.internmc.facebook.com/intern/diff/D112906313/) --- .../graph/ops/glsl/sdpa_compute_out_coop.yaml | 8 ++++ .../vulkan/runtime/graph/ops/impl/SDPA.cpp | 48 ++++++++++++++++--- backends/vulkan/runtime/graph/ops/impl/SDPA.h | 33 +++++++++---- .../vulkan/test/custom_ops/impl/TestSDPA.cpp | 23 ++++++--- backends/vulkan/test/custom_ops/test_sdpa.cpp | 21 +++++++- 5 files changed, 109 insertions(+), 24 deletions(-) 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"));