Add benchmarks for transposed vs standard KV cache layout

Benchmarks comparing transposed [B, H, S, D] vs standard [B, S, H, D] KV
cache layouts in custom_sdpa and update_cache ops using Llama 3 8B config
(32 Q heads, 8 KV heads, D=128). Both C++ (Google Benchmark) and Python
benchmarks are included, covering decode (seq_len=1) at various cache fill
levels and prefill scenarios.

Results on Apple M-series show transposed cache significantly improves SDPA
performance at longer cache fills (1.64x at start_pos=1024, 1.13x for
prefill seq_len=512) due to better memory locality in the attn_score @ V
GEMM — V stride along S_kv changes from H*D to D.

Authored with Claude.

Differential Revision: [D99677680](https://our.internmc.facebook.com/intern/diff/D99677680/)

[ghstack-poisoned]

Author: kimishpatel

extension/llm/custom_ops/BUCK

@@ -105,6 +105,21 @@ fbcode_target(_kind = runtime.python_test,
     ],
 )
 
+fbcode_target(_kind = runtime.python_binary,
+    name = "bench_transposed_cache_py",
+    srcs = [
+        "bench_transposed_cache.py",
+    ],
+    main_function = "executorch.extension.llm.custom_ops.bench_transposed_cache.main",
+    preload_deps = [
+        ":custom_ops_aot_lib_mkl_noomp",
+        ":custom_ops_aot_py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
 fbcode_target(_kind = cpp_benchmark,
     name = "bench_sdpa",
     srcs = ["bench_sdpa.cpp"],
@@ -117,3 +132,14 @@ fbcode_target(_kind = cpp_benchmark,
         "//executorch/runtime/core/exec_aten/testing_util:tensor_util",
     ],
 )
+
+cpp_benchmark(
+    name = "bench_transposed_cache",
+    srcs = ["bench_transposed_cache.cpp"],
+    deps = [
+        "fbsource//third-party/benchmark:benchmark",
+        "//executorch/extension/llm/custom_ops:custom_ops_mkl_noomp",
+        "//executorch/runtime/core/exec_aten:lib",
+        "//executorch/runtime/core/exec_aten/testing_util:tensor_util",
+    ],
+)