Update on "[Executorch] Add non-flash SDPA for decode"

Add cpu_sdpa template function in op_sdpa_impl.h that provides a
simpler SDPA implementation using standard GEMM (no tiling). This is
useful as a baseline and for cases where flash attention is not optimal.

The implementation uses a single SeqDim parameter for all tensors and
supports causal masking, attention masks, GQA, and multi-threading.

During decode (seq_len == 1), the tiled flash attention implementation
has unnecessary overhead from its blocking/tiling logic. The simpler
unfused SDPA path using direct GEMM is more efficient for single-query
attention, yielding ~25-30% decode throughput improvement on S25
(41 -> 53 tok/s for 1.4B parameter model).

This makes cpu_sdpa always available (previously gated behind
ET_USE_UNFUSED_SDPA) and dispatches to it when seq_len == 1 and
inputs are not quantized. Prefill continues to use flash attention.

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

[ghstack-poisoned]

Author: kimishpatel

.ci/scripts/test_model_e2e.sh

@@ -354,7 +354,7 @@ EOF
     fi
     ;;
   qwen3_5_moe)
-    RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 32"
+    RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0"
     ;;
   voxtral_realtime)
     RUNNER_ARGS="--model_path ${MODEL_DIR}/model.pte --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --preprocessor_path ${MODEL_DIR}/$PREPROCESSOR --audio_path ${MODEL_DIR}/$AUDIO_FILE --temperature 0"

.github/workflows/_test_cadence.yml

@@ -0,0 +1,77 @@
+name: Test Cadence
+
+permissions:
+  id-token: write
+  contents: read
+
+on:
+  workflow_call:
+    inputs:
+      docker-image:
+        description: 'Docker image to use'
+        required: false
+        type: string
+        default: ci-image:executorch-ubuntu-22.04-clang12
+      runner:
+        description: 'Runner type'
+        required: false
+        type: string
+        default: linux.8xlarge.memory
+      ref:
+        description: 'Git ref to checkout'
+        required: false
+        type: string
+        default: ${{ github.sha }}
+      timeout:
+        description: 'Job timeout in minutes'
+        required: false
+        type: number
+        default: 90
+
+jobs:
+  test-aot:
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    with:
+      job-name: test-aot
+      runner: ${{ inputs.runner }}
+      docker-image: ${{ inputs.docker-image }}
+      submodules: recursive
+      ref: ${{ inputs.ref }}
+      timeout: ${{ inputs.timeout }}
+      script: |
+        set -eux
+        conda create -y -n cadence_test python=3.12 > /dev/null
+        conda activate cadence_test
+
+        ./install_requirements.sh > /dev/null
+        pip install -e . --no-build-isolation > /dev/null
+        pip install beartype later pyre_extensions pytest-xdist
+
+        python -m pytest backends/cadence/aot/tests/ -v -n auto
+
+  test-ops:
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    with:
+      job-name: test-ops
+      runner: ${{ inputs.runner }}
+      docker-image: ${{ inputs.docker-image }}
+      submodules: recursive
+      ref: ${{ inputs.ref }}
+      timeout: ${{ inputs.timeout }}
+      download-artifact: cadence-runner-build
+      script: |
+        set -eux
+        conda create -y -n cadence_test python=3.12 > /dev/null
+        conda activate cadence_test
+
+        ./install_requirements.sh > /dev/null
+        pip install -e . --no-build-isolation > /dev/null
+        pip install beartype later pyre_extensions pytest-xdist
+
+        # Use the pre-built runner from the build job
+        mkdir -p cmake-out/backends/cadence
+        cp "${RUNNER_ARTIFACT_DIR}/cadence_runner" cmake-out/backends/cadence/cadence_runner
+        chmod +x cmake-out/backends/cadence/cadence_runner
+
+        export PYTHONPATH="${PYTHONPATH:-}:$(pwd)/backends/cadence/utils/FACTO"
+        python -m pytest examples/cadence/operators/ -v -n auto

.github/workflows/build-cadence-runner.yml

@@ -1,4 +1,4 @@
-name: Build Cadence
+name: Cadence Build & Test
 
 on:
   pull_request:
@@ -13,7 +13,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  cpu:
+  cpu-build:
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     permissions:
       id-token: write
@@ -25,6 +25,7 @@ jobs:
       submodules: recursive
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 90
+      upload-artifact: cadence-runner-build
       script: |
         set -eux
         # The generic Linux job chooses to use base env, not the one setup by the image
@@ -33,3 +34,15 @@ jobs:
 
         ./install_requirements.sh > /dev/null
         bash backends/cadence/build_cadence_runner.sh
+
+        # Copy runner binary to artifact dir for downstream test jobs
+        cp cmake-out/backends/cadence/cadence_runner "${RUNNER_ARTIFACT_DIR}/"
+
+  cpu-test:
+    needs: cpu-build
+    permissions:
+      id-token: write
+      contents: read
+    uses: ./.github/workflows/_test_cadence.yml
+    with:
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}

.github/workflows/cuda.yml

@@ -145,8 +145,8 @@ jobs:
         # Run CUDA backend Python tests
         python -m pytest backends/cuda/tests backends/cuda/passes/tests -v -o "addopts="
 
-        # Run quantize roundtrip tests (Qwen 3.5 MoE save/load prequantized)
-        python -m pytest examples/models/qwen3_5_moe/test_quantize_roundtrip.py -v -o "addopts="
+        # Run Qwen 3.5 MoE tests (quantize roundtrip + TurboQuant KV cache)
+        python -m pytest examples/models/qwen3_5_moe/test_quantize_roundtrip.py examples/models/qwen3_5_moe/test_turboquant.py -v -o "addopts="
 
   export-model-cuda-artifact:
     name: export-model-cuda-artifact

backends/aoti/common_shims_slim.cpp

@@ -134,6 +134,10 @@ int32_t aoti_torch_dtype_int8() {
   return 1; // ScalarType::Char
 }
 
+int32_t aoti_torch_dtype_uint8() {
+  return 0; // ScalarType::Byte
+}
+
 int32_t aoti_torch_dtype_bool() {
   return 11; // ScalarType::Bool
 }

backends/aoti/common_shims_slim.h

@@ -76,6 +76,7 @@ AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int64();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int32();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int16();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int8();
+AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_uint8();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bool();
 
 // ============================================================

backends/aoti/slim/c10/core/ScalarType.h

@@ -23,7 +23,7 @@ using BFloat16 = ::executorch::runtime::etensor::BFloat16;
 /// Enum representing the scalar type (dtype) of tensor elements.
 /// Note: Enum values must match PyTorch's c10::ScalarType for compatibility.
 enum class ScalarType : int8_t {
-  // Byte = 0,     // uint8_t - not currently needed
+  Byte = 0, // uint8_t
   Char = 1, // int8_t
   Short = 2, // int16_t
   Int = 3, // int32_t
@@ -43,6 +43,7 @@ enum class ScalarType : int8_t {
 };
 
 // Type alias constants for convenience
+constexpr ScalarType kByte = ScalarType::Byte;
 constexpr ScalarType kChar = ScalarType::Char;
 constexpr ScalarType kShort = ScalarType::Short;
 constexpr ScalarType kInt = ScalarType::Int;
@@ -56,6 +57,8 @@ constexpr ScalarType kBFloat16 = ScalarType::BFloat16;
 /// @return The size in bytes of a single element.
 inline size_t elementSize(ScalarType t) {
   switch (t) {
+    case ScalarType::Byte:
+      return sizeof(uint8_t);
     case ScalarType::Char:
       return sizeof(int8_t);
     case ScalarType::Short:
@@ -80,6 +83,8 @@ inline size_t elementSize(ScalarType t) {
 /// @return The name of the scalar type.
 inline const char* toString(ScalarType t) {
   switch (t) {
+    case ScalarType::Byte:
+      return "Byte";
     case ScalarType::Char:
       return "Char";
     case ScalarType::Short:
@@ -114,6 +119,7 @@ inline bool isFloatingType(ScalarType t) {
 /// @return true if the scalar type is integral, false otherwise.
 inline bool isIntegralType(ScalarType t, bool includeBool) {
   switch (t) {
+    case ScalarType::Byte:
     case ScalarType::Char:
     case ScalarType::Short:
     case ScalarType::Int:
@@ -138,6 +144,7 @@ inline bool isBoolType(ScalarType t) {
 /// @return true if the scalar type is valid, false otherwise.
 inline bool isValidScalarType(ScalarType t) {
   switch (t) {
+    case ScalarType::Byte:
     case ScalarType::Char:
     case ScalarType::Short:
     case ScalarType::Int:

backends/arm/runtime/VGFBackend.cpp

@@ -88,7 +88,14 @@ void vkml_free_basics(
 
 class VGFBackend final : public ::executorch::runtime::BackendInterface {
  public:
-  VGFBackend() {
+  VGFBackend() = default;
+
+  // Lazy Vulkan init — runs on first use, not in the constructor.
+  void ensure_initialized() {
+    if (is_initialized_) {
+      return;
+    }
+
     VkResult result;
 
     // Fetch basic vulkan objects once
@@ -122,6 +129,7 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface {
 
   bool is_available() const override {
     ET_LOG(Info, "Checking VGFBackend is available");
+    const_cast<VGFBackend*>(this)->ensure_initialized();
     if (!is_initialized_) {
       return false;
     }
@@ -134,6 +142,7 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface {
       ArrayRef<CompileSpec> compile_specs) const override {
     ET_LOG(Info, "Entered VGF init");
 
+    const_cast<VGFBackend*>(this)->ensure_initialized();
     if (!is_initialized_) {
       ET_LOG(
           Error,
@@ -334,6 +343,16 @@ VkResult vkml_allocate_basics(
   }
   volkLoadInstance(*instance);
 
+  // Bail out if the driver lacks ARM tensor/datagraph extensions.
+  if (!vkCreateTensorARM) {
+    ET_LOG(
+        Error,
+        "Vulkan driver does not support ARM tensor extensions (VK_ARM_tensors)");
+    vkDestroyInstance(*instance, nullptr);
+    *instance = VK_NULL_HANDLE;
+    return VK_ERROR_FEATURE_NOT_PRESENT;
+  }
+
   // Pick first GPU
   uint32_t gpu_count = 0;
   vkEnumeratePhysicalDevices(*instance, &gpu_count, nullptr);

backends/arm/runtime/VGFSetup.cpp

@@ -538,7 +538,7 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs) {
       .pNext = nullptr,
       .flags = 0,
       .bindingCount = static_cast<uint32_t>(layout_bindings.size()),
-      layout_bindings.data(),
+      .pBindings = layout_bindings.data(),
   };
   result =
       vkCreateDescriptorSetLayout(vk_device, &layout_info, nullptr, &vk_layout);

backends/arm/runtime/targets.bzl

@@ -32,3 +32,35 @@ def define_common_targets():
             "fbsource//third-party/ethos-u-core-driver:core_driver",
         ],
     )
+    runtime.cxx_library(
+        name = "vgf_backend",
+        srcs = [
+            "VGFBackend.cpp",
+            "VGFSetup.cpp",
+            # Volk must be compiled directly into this target so its global
+            # function-pointer variables live in the same linkage unit.
+            # Linking from a separate static library causes the linker to
+            # drop the symbols when building a shared library.
+            "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:volk_arm_src",
+        ],
+        exported_headers = ["VGFSetup.h"],
+        # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole)
+        link_whole = True,
+        supports_python_dlopen = True,
+        compiler_flags = [
+            "-Wno-global-constructors",
+            "-Wno-header-hygiene",
+            "-Wno-unused-variable",
+            "-Wno-missing-field-initializers",
+            "-DUSE_VULKAN_WRAPPER",
+            "-DUSE_VULKAN_VOLK",
+        ],
+        visibility = ["PUBLIC"],
+        deps = [
+            "//executorch/runtime/backend:interface",
+            "//executorch/runtime/core:core",
+            "fbsource//third-party/arm-vgf-library/v0.8.0/src:vgf",
+            "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:volk_arm",
+            "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:vulkan-headers",
+        ],
+    )