From 5d2c19ed11b345f7a3df8ee708af435eed9f331b Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Tue, 21 Jul 2026 08:23:49 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../vulkan/test/custom_ops/impl/TestSDPA.cpp | 72 ++++ backends/vulkan/test/custom_ops/targets.bzl | 1 + backends/vulkan/test/custom_ops/test_sdpa.cpp | 341 ++++++++++++++++++ 3 files changed, 414 insertions(+) create mode 100644 backends/vulkan/test/custom_ops/impl/TestSDPA.cpp create mode 100644 backends/vulkan/test/custom_ops/test_sdpa.cpp 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..dc908de55e6 --- /dev/null +++ b/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 + +#include + +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& 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(-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); + + 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 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..15040c781db --- /dev/null +++ b/backends/vulkan/test/custom_ops/test_sdpa.cpp @@ -0,0 +1,341 @@ +// 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 + std::string impl_selector = "default"; +}; + +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) { + 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 + "]"; + + 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(config.impl_selector); + + // 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; + + 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"; + test_cases.push_back(create_sdpa_test_case(cfg, dtype, storage)); + } + 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)); + } + } + + if (decode_only_mode()) { + 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. + { + SDPAConfig dec{64, 8, 2, 1, 32, "accu", "decode"}; + test_cases.push_back( + create_sdpa_test_case(dec, vkapi::kFloat, utils::kTexture3D)); + SDPAConfig pre{64, 8, 2, 16, 16, "accu", "prefill"}; + test_cases.push_back( + create_sdpa_test_case(pre, vkapi::kFloat, utils::kTexture3D)); + } + + 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; +}