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..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 @@ -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,13 @@ sdpa_compute_out_coop: - VALUE: half shader_variants: - 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 3efb834725d..c4259d4c530 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,73 @@ 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): 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, + 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) != + 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, + "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, @@ -297,8 +352,36 @@ 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"; + // 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"; + } + } 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 +402,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 +410,24 @@ 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); + // 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)}; } @@ -498,7 +598,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 +612,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 +642,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 +718,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 +754,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..fa51ebef111 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.h @@ -0,0 +1,83 @@ +/* + * 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); + +// 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. +// 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, + 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 new file mode 100644 index 00000000000..4a59210cf9c --- /dev/null +++ b/backends/vulkan/test/custom_ops/impl/TestSDPA.cpp @@ -0,0 +1,175 @@ +/* + * 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 + +#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` +// symint, but input_pos is not a free parameter: the op enforces +// 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. +// +// 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 (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 +void test_sdpa(ComputeGraph& graph, const std::vector& args) { + int arg_idx = 0; + const ValueRef q = args.at(arg_idx++); + const ValueRef k_cache = args.at(arg_idx++); + const ValueRef v_cache = args.at(arg_idx++); + const ValueRef impl_selector_str = args.at(arg_idx++); + const ValueRef out = args.at(arg_idx++); + + 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, 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); + VK_CHECK_COND(context_len >= seq_len); + const int32_t input_pos_val = + utils::safe_downcast(context_len - seq_len); + + const ValueRef input_pos_symint = graph.add_symint(input_pos_val); + + // 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(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(kShaderOverrideForceNonGqa); + } + + test_sdpa_impl( + graph, q, k_cache, v_cache, input_pos_symint, out, shader_override); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(test_etvk.test_sdpa.default, test_sdpa); +} + +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index 7ff7b6ec426..5d1045173a2 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -108,3 +108,4 @@ def define_common_targets(is_fbcode = False): define_custom_op_test_binary("test_conv1d_pw") define_custom_op_test_binary("test_conv1d_dw") define_custom_op_test_binary("test_fpa_q4gsw_linear") + define_custom_op_test_binary("test_sdpa") diff --git a/backends/vulkan/test/custom_ops/test_sdpa.cpp b/backends/vulkan/test/custom_ops/test_sdpa.cpp new file mode 100644 index 00000000000..55d2173e3cd --- /dev/null +++ b/backends/vulkan/test/custom_ops/test_sdpa.cpp @@ -0,0 +1,379 @@ +// 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 +#include +#include +#include +#include + +#include +#include + +#include "utils.h" + +using namespace executorch::vulkan::prototyping; +using namespace vkcompute; + +// Correctness is only checked for these small shapes; larger perf shapes throw +// std::invalid_argument from the reference (framework marks them SKIPPED) to +// avoid an O(H*S*context*D) CPU reference on the large perf matrix. +static constexpr int64_t kRefContextLenLimit = 256; + +// When true (env SDPA_NO_CHAIN=1), each graph.execute() dispatches the SDPA op +// exactly once (op_invocations_per_execute=1), disabling the framework's +// probe-then-scale chaining. Chaining stacks many back-to-back QK/softmax/AV +// triples in one command buffer; consecutive triples pipeline on the GPU so +// their timestamp windows overlap, which inflates/misattributes per-dispatch +// durations (notably the AV dispatch) at the large chained_dispatches factors +// picked for the cheap small-context decode cases. One-invocation-per-execute +// removes the inter-invocation overlap; stability comes from warmup + a large +// median-of-N instead. +static bool no_chain_mode() { + const char* v = std::getenv("SDPA_NO_CHAIN"); + return v != nullptr && v[0] == '1'; +} + +// When true (env SDPA_DECODE_ONLY=1), generate only the 9 decode cases. +static bool decode_only_mode() { + const char* v = std::getenv("SDPA_DECODE_ONLY"); + return v != nullptr && v[0] == '1'; +} + +// LLM SDPA (llama.custom_sdpa) shape: +// q: [1, S, n_heads, head_dim] (DHSB, width-packed) +// k/v cache:[1, context_len, n_kv_heads, head_dim] +struct SDPAConfig { + int64_t head_dim; + int64_t n_heads; + int64_t n_kv_heads; + int64_t seq_len; // S: query tokens (1 for decode, >1 for prefill) + int64_t context_len; // total KV length (kv_len) + std::string model; // label only + std::string regime; // "decode" / "prefill", label only +}; + +static std::vector as_float_data(const ValueSpec& spec) { + if (spec.dtype == vkapi::kFloat) { + return spec.get_float_data(); + } + if (spec.dtype == vkapi::kHalf) { + const auto& half_bits = spec.get_half_data(); + std::vector out(half_bits.size()); + for (size_t i = 0; i < half_bits.size(); ++i) { + out[i] = half_to_float(half_bits[i]); + } + return out; + } + throw std::invalid_argument("as_float_data: unsupported dtype"); +} + +static TestCase create_sdpa_test_case( + const SDPAConfig& config, + vkapi::ScalarType dtype, + utils::StorageType storage_type, + const std::string& impl) { + TestCase test_case; + + const bool is_perf = config.context_len > kRefContextLenLimit; + const std::string prefix = is_perf ? "PERF" : "ACCU"; + const std::string storage_str = repr_str(storage_type, utils::kWidthPacked); + const std::string dtype_str = dtype_short(dtype); + + const std::string shape = "D" + std::to_string(config.head_dim) + " H" + + std::to_string(config.n_heads) + " Hkv" + + std::to_string(config.n_kv_heads) + " S" + + std::to_string(config.seq_len) + " C" + + std::to_string(config.context_len); + + 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)); + test_case.set_operator_name("test_etvk.test_sdpa.default"); + + // q: [1, S, n_heads, head_dim] + ValueSpec q( + {1, config.seq_len, config.n_heads, config.head_dim}, + dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + + // k_cache / v_cache: [1, context_len, n_kv_heads, head_dim] + ValueSpec k_cache( + {1, config.context_len, config.n_kv_heads, config.head_dim}, + dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + ValueSpec v_cache( + {1, config.context_len, config.n_kv_heads, config.head_dim}, + dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + + ValueSpec impl_selector = ValueSpec::make_string(impl); + + // out: [1, S, n_heads, head_dim] + ValueSpec output( + {1, config.seq_len, config.n_heads, config.head_dim}, + dtype, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + + test_case.add_input_spec(q); + test_case.add_input_spec(k_cache); + test_case.add_input_spec(v_cache); + test_case.add_input_spec(impl_selector); + test_case.add_output_spec(output); + + if (no_chain_mode()) { + test_case.set_op_invocations_per_execute(1); + } + + if (dtype == vkapi::kHalf) { + test_case.set_abs_tolerance(1e-2f); + test_case.set_rel_tolerance(1e-2f); + } else { + test_case.set_abs_tolerance(1e-3f); + test_case.set_rel_tolerance(1e-3f); + } + + return test_case; +} + +// Reference: causal SDPA over the KV cache. +// q:[1,S,H,D], k/v cache:[1,C,Hkv,D], input_pos = C - S. +// For query row s (absolute position input_pos + s), attends to cache +// positions [0, input_pos + s]. GQA: head h maps to kv head h / (H/Hkv). +static void sdpa_reference_impl(TestCase& test_case) { + const auto& q = test_case.inputs()[0]; + const auto& k = test_case.inputs()[1]; + const auto& v = test_case.inputs()[2]; + + const auto q_sizes = q.get_tensor_sizes(); + const auto k_sizes = k.get_tensor_sizes(); + + const int64_t S = q_sizes[1]; + const int64_t H = q_sizes[2]; + const int64_t D = q_sizes[3]; + const int64_t C = k_sizes[1]; + const int64_t Hkv = k_sizes[2]; + + if (C > kRefContextLenLimit) { + throw std::invalid_argument("sdpa reference: perf shape, skipping"); + } + + const int64_t input_pos = C - S; + const int64_t heads_per_kv = H / Hkv; + const float scale = 1.0f / std::sqrt(static_cast(D)); + + const auto q_data = as_float_data(q); + const auto k_data = as_float_data(k); + const auto v_data = as_float_data(v); + + ValueSpec& output = test_case.outputs()[0]; + auto& ref = output.get_ref_float_data(); + ref.assign(S * H * D, 0.0f); + + // Index helpers (contiguous WHCN-flattened as [1, dim1, dim2, dim3]). + auto q_idx = [&](int64_t s, int64_t h, int64_t d) { + return (s * H + h) * D + d; + }; + auto kv_idx = [&](int64_t c, int64_t hk, int64_t d) { + return (c * Hkv + hk) * D + d; + }; + + std::vector scores(C); + for (int64_t s = 0; s < S; ++s) { + const int64_t attend_len = input_pos + s + 1; // causal + for (int64_t h = 0; h < H; ++h) { + const int64_t hk = h / heads_per_kv; + + float max_score = -std::numeric_limits::infinity(); + for (int64_t c = 0; c < attend_len; ++c) { + float dot = 0.0f; + for (int64_t d = 0; d < D; ++d) { + dot += q_data[q_idx(s, h, d)] * k_data[kv_idx(c, hk, d)]; + } + dot *= scale; + scores[c] = dot; + max_score = std::max(max_score, dot); + } + + float denom = 0.0f; + for (int64_t c = 0; c < attend_len; ++c) { + scores[c] = std::exp(scores[c] - max_score); + denom += scores[c]; + } + + for (int64_t d = 0; d < D; ++d) { + float acc = 0.0f; + for (int64_t c = 0; c < attend_len; ++c) { + acc += scores[c] * v_data[kv_idx(c, hk, d)]; + } + ref[q_idx(s, h, d)] = acc / denom; + } + } + } +} + +// FLOPs: QK (2*S*C*D) + AV (2*S*C*D) per head, summed over heads. Softmax +// is negligible. Uses the causal-average context (~C/2) is ignored; report +// full-C dense FLOPs as an upper bound proxy. +static int64_t sdpa_flop_calculator(const TestCase& test_case) { + const auto q_sizes = test_case.inputs()[0].get_tensor_sizes(); + const auto k_sizes = test_case.inputs()[1].get_tensor_sizes(); + const int64_t S = q_sizes[1]; + const int64_t H = q_sizes[2]; + const int64_t D = q_sizes[3]; + const int64_t C = k_sizes[1]; + return 4 * H * S * C * D; +} + +static std::vector generate_sdpa_test_cases() { + std::vector test_cases; + + struct ModelDims { + std::string name; + int64_t head_dim; + int64_t n_heads; + int64_t n_kv_heads; + }; + + const std::vector models = { + {"Llama-3.2-1B", 64, 32, 8}, + {"Qwen3-0.6B", 128, 16, 8}, + {"Phi-4-mini", 128, 24, 8}, + }; + + // Decode: S=1, sweep context_len. + const std::vector decode_context_lens = {512, 1024, 4096}; + // Prefill: S == context_len. + const std::vector prefill_seq_lens = {128, 512}; + + // Perf runs use fp16 texture (matches the LLM decode/prefill production + // path). A couple of small ACCU shapes validate correctness in fp32. + 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; + cfg.head_dim = m.head_dim; + cfg.n_heads = m.n_heads; + cfg.n_kv_heads = m.n_kv_heads; + cfg.seq_len = 1; + cfg.context_len = c; + cfg.model = m.name; + cfg.regime = "decode"; + for (const auto& impl : decode_impls) { + test_cases.push_back(create_sdpa_test_case(cfg, dtype, storage, impl)); + } + } + if (decode_only_mode()) { + continue; + } + for (int64_t s : prefill_seq_lens) { + SDPAConfig cfg; + cfg.head_dim = m.head_dim; + cfg.n_heads = m.n_heads; + cfg.n_kv_heads = m.n_kv_heads; + cfg.seq_len = s; + cfg.context_len = s; + cfg.model = m.name; + cfg.regime = "prefill"; + test_cases.push_back( + create_sdpa_test_case(cfg, dtype, storage, "default")); + } + } + + if (decode_only_mode()) { + return test_cases; + } + + // 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). + { + // 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"}, + }; + 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)); + } + } + } + + // 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")); + } + + return test_cases; +} + +int main(int argc, char* argv[]) { + (void)argc; + (void)argv; + + set_debugging(false); + set_print_output(false); + set_print_latencies(false); + set_use_gpu_timestamps(true); + + print_performance_header(); + std::cout << "SDPA (llama.custom_sdpa) Benchmark" << std::endl; + print_separator(); + + ReferenceComputeFunc ref_fn = sdpa_reference_impl; + + const bool decode_only = decode_only_mode(); + const int warmup_runs = decode_only ? 10 : 3; + const int benchmark_runs = decode_only ? 30 : 10; + + auto results = execute_test_cases( + generate_sdpa_test_cases, + sdpa_flop_calculator, + "SDPA", + warmup_runs, + benchmark_runs, + ref_fn); + + return 0; +} 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); +}