[ET-VK][sdpa] Add SDPA operator perf benchmark binary (test_sdpa)

Differential Revision: D112906312

Pull Request resolved: #21062

Author: SS-JIA

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

@@ -0,0 +1,72 @@
+/*
+ * 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 <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+
+namespace vkcompute {
+
+// 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;
+// exposing input_pos as its own ValueSpec would add a redundant degree of
+// freedom that a test could set inconsistent with the cache shape.
+//
+// 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.
+//
+// Args: q, k_cache, v_cache, impl_selector, out
+void test_sdpa(ComputeGraph& graph, const std::vector<ValueRef>& 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",
+      "test_sdpa only supports the 'default' impl_selector");
+
+  const int64_t seq_len = graph.size_at<int64_t>(-3, q);
+  const int64_t context_len = graph.size_at<int64_t>(-3, k_cache);
+  VK_CHECK_COND(context_len >= seq_len);
+  const int32_t input_pos_val =
+      utils::safe_downcast<int32_t>(context_len - seq_len);
+
+  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,
+   });
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(test_etvk.test_sdpa.default, test_sdpa);
+}
+
+} // namespace vkcompute

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")