cuda_backend.cpp

/*
 * 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 <c10/util/safe_numerics.h>
#include <executorch/extension/cuda/runtime_api.h>
#include <executorch/runtime/backend/interface.h>
#include <executorch/runtime/backend/options.h>
#include <executorch/runtime/core/error.h>
#include <executorch/runtime/core/evalue.h>
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
#include <cctype>
#include <cstdio>

#include <array>
#include <atomic>
#include <filesystem>
#include <fstream>
#include <mutex>
#include <optional>
#include <string>
#include <string_view>
#include <unordered_map>
#include <unordered_set>
#include <vector>

// Include SlimTensor headers for CUDA backend
#include <executorch/backends/aoti/slim/c10/core/Device.h>
#include <executorch/backends/aoti/slim/c10/cuda/Exception.h>
#include <executorch/backends/aoti/slim/core/slim_tensor.h>
#include <executorch/backends/aoti/slim/core/storage.h>
#include <executorch/backends/aoti/slim/cuda/guard.h>
#include <executorch/backends/aoti/slim/factory/empty.h>
#include <executorch/backends/aoti/slim/factory/from_blob.h>
#include <executorch/backends/aoti/slim/util/array_ref_util.h>
#include <executorch/extension/cuda/caller_stream.h>

// Include our shim layer headers
#include <executorch/backends/aoti/aoti_delegate_handle.h>
#include <executorch/backends/aoti/utils.h>
#include <executorch/backends/cuda/runtime/cuda_allocator.h>
#include <executorch/backends/cuda/runtime/cuda_delegate_handle.h>
#include <executorch/backends/cuda/runtime/cuda_kv_cache.h>
#include <executorch/backends/cuda/runtime/cuda_mutable_state.h>
#include <executorch/backends/cuda/runtime/cuda_weight_cache.h>
#include <executorch/backends/cuda/runtime/platform/platform.h>
#include <executorch/backends/cuda/runtime/shims/memory.h>
#include <executorch/backends/cuda/runtime/utils.h>
#include <executorch/extension/llm/cache/cache_registry.h>

namespace executorch::backends::cuda {

using namespace std;
using namespace aoti;

using executorch::runtime::ArrayRef;
using executorch::runtime::BackendExecutionContext;
using executorch::runtime::BackendInitContext;
using executorch::runtime::BackendOption;
using executorch::runtime::BackendOptionContext;
using executorch::runtime::CompileSpec;
using executorch::runtime::DelegateHandle;
using executorch::runtime::Error;
using executorch::runtime::EValue;
using executorch::runtime::FreeableBuffer;
using executorch::runtime::kMaxOptionValueLength;
using executorch::runtime::NamedDataMap;
using executorch::runtime::Result;
using executorch::runtime::Span;
using executorch::runtime::etensor::Tensor;

// SlimTensor type aliases
using cuda::CudaGraphPhase;
using slim::SlimTensor;
using slim::c10::Device;

namespace {
constexpr char kSkipCopyOutputToCpuForMethod[] =
    "skip_copy_output_to_cpu_for_method";
constexpr char kUseSharedCudaStream[] = "use_shared_cuda_stream";

constexpr char kEnableCudaGraphForMethod[] = "enable_cuda_graph_for_method";
constexpr int kCudaGraphWarmupSteps = 3;
constexpr char kWeightSharingAcrossMethods[] = "weight_sharing_across_methods";
} // anonymous namespace

// Advances the off-graph KV cache past a step's tokens. Every successful exit
// of execute() ends here, and only those: a run that failed never wrote its
// tokens, and committing them would desynchronise the cache's length from the
// history the kernels actually hold.
Error commit_offgraph_kv_step(CudaDelegateHandle* handle, int64_t width) {
  if (handle->kv_cache == nullptr) {
    return Error::Ok;
  }
  return handle->kv_cache->commit_step(width);
}

class ET_EXPERIMENTAL CudaBackend final
    : public ::executorch::runtime::BackendInterface {
 private:
  // Trim leading/trailing whitespace from a view of the string.
  static std::string_view trim(std::string_view s) {
    size_t start = 0;
    while (start < s.size() &&
           std::isspace(static_cast<unsigned char>(s[start]))) {
      ++start;
    }
    size_t end = s.size();
    while (end > start &&
           std::isspace(static_cast<unsigned char>(s[end - 1]))) {
      --end;
    }
    return s.substr(start, end - start);
  }

  // Check if method_name appears in a comma-separated list.
  static bool method_in_csv(
      const std::string& method_name,
      const std::string& csv) {
    size_t pos = 0;
    while (pos <= csv.size()) {
      const size_t comma = csv.find(',', pos);
      const std::string_view token =
          trim(std::string_view(csv).substr(pos, comma - pos));
      if (!token.empty() && token == method_name) {
        return true;
      }
      if (comma == std::string::npos) {
        break;
      }
      pos = comma + 1;
    }
    return false;
  }

  void set_cuda_graph_method(
      const std::array<char, kMaxOptionValueLength>& raw) {
    std::lock_guard<std::mutex> guard(cuda_graph_method_mutex_);
    cuda_graph_method_ = std::string(raw.data());
  }

  bool should_use_cuda_graph_for_method(const std::string& method_name) const {
    if (method_name.empty()) {
      return false;
    }
    std::lock_guard<std::mutex> guard(cuda_graph_method_mutex_);
    return method_in_csv(method_name, cuda_graph_method_);
  }

  // Enable the legacy dense-blob per-FQN cache. New FQN artifacts use
  // their FQN-addressed data keys automatically.
  void set_weight_sharing_across_methods(bool enabled) {
    weight_sharing_across_methods_.store(enabled, std::memory_order_relaxed);
  }

  bool is_weight_sharing_across_methods_enabled() const {
    return weight_sharing_across_methods_.load(std::memory_order_relaxed);
  }

  Error load_function_pointers_into_handle(
      void* so_handle,
      cuda::CudaDelegateHandle* handle) const {
#define LOAD_SYMBOL(member, name)                                    \
  do {                                                               \
    auto symbol_res = get_function(so_handle, #name);                \
    if (!symbol_res.ok()) {                                          \
      return symbol_res.error();                                     \
    }                                                                \
    handle->member = reinterpret_cast<name##Func>(symbol_res.get()); \
  } while (0)

    LOAD_SYMBOL(create_with_device, AOTInductorModelContainerCreateWithDevice);

    LOAD_SYMBOL(delete_container, AOTInductorModelContainerDelete);

    LOAD_SYMBOL(get_num_inputs, AOTInductorModelContainerGetNumInputs);

    LOAD_SYMBOL(get_num_outputs, AOTInductorModelContainerGetNumOutputs);

    LOAD_SYMBOL(run, AOTInductorModelContainerRun);
#undef LOAD_SYMBOL

    auto symbol_res =
        get_function(so_handle, "AOTInductorModelUpdateConstantsFromBlob");
    if (symbol_res.ok()) {
      handle->update_constants_from_blob =
          reinterpret_cast<AOTInductorModelUpdateConstantsFromBlobFunc>(
              symbol_res.get());
    } else {
      ET_LOG(
          Info,
          "Failed to load AOTInductorModelUpdateConstantsFromBlob. This .so is probably compiled on an old version of torch (<2.9.0)");
    }

    // Load constant management symbols (optional — needed for cross-method
    // buffer sharing). These are available in torch >= 2.6.
#define LOAD_OPTIONAL_SYMBOL(member, name)                            \
  do {                                                                \
    auto res = get_function(so_handle, #name);                        \
    handle->member =                                                  \
        res.ok() ? reinterpret_cast<name##Func>(res.get()) : nullptr; \
  } while (0)

    LOAD_OPTIONAL_SYMBOL(
        get_num_constants, AOTInductorModelContainerGetNumConstants);
    LOAD_OPTIONAL_SYMBOL(
        get_constant_name, AOTInductorModelContainerGetConstantName);
    LOAD_OPTIONAL_SYMBOL(
        get_constant_original_fqn,
        AOTInductorModelContainerGetConstantOriginalFQN);
    LOAD_OPTIONAL_SYMBOL(
        get_constant_dtype, AOTInductorModelContainerGetConstantDtype);
    LOAD_OPTIONAL_SYMBOL(
        get_constant_data_size, AOTInductorModelContainerGetConstantDataSize);
    LOAD_OPTIONAL_SYMBOL(
        extract_constants_map, AOTInductorModelContainerExtractConstantsMap);
    LOAD_OPTIONAL_SYMBOL(
        update_user_managed_constant_buffer_pairs,
        AOTInductorModelContainerUpdateUserManagedConstantBufferPairs);
#undef LOAD_OPTIONAL_SYMBOL

    return Error::Ok;
  }

 public:
  bool is_available() const override {
    return 1;
  }

  Error set_option(
      ET_UNUSED BackendOptionContext& context,
      const executorch::runtime::Span<BackendOption>& backend_options)
      override {
    for (const auto& option : backend_options) {
      if (std::strcmp(option.key, kSkipCopyOutputToCpuForMethod) == 0) {
        // Deprecated, no-op option. CUDA delegate IO is now GPU-resident under
        // device memory planning, and host<->device transfers are handled by
        // graph-level et_copy ops. To skip those copies, export the .pte with
        // ExecutorchBackendConfig.propagate_device_config =
        // PropagateDeviceConfig( skip_d2h_for_method_outputs=...,
        // skip_h2d_for_method_inputs=...).
        ET_LOG(
            Info,
            "Runtime backend option '%s' is DEPRECATED and no longer has any "
            "effect; ignoring it.",
            kSkipCopyOutputToCpuForMethod);
      } else if (std::strcmp(option.key, kUseSharedCudaStream) == 0) {
        // Refused rather than ignored: this option was the only thing ordering
        // methods driven from different threads, so silently dropping it would
        // give such a caller unordered device work and wrong results. Methods
        // now run on the calling thread's stream, which orders methods called
        // from one thread but not across threads; a caller that needs that must
        // order the calls itself.
        if (auto* val = std::get_if<bool>(&option.value)) {
          if (*val) {
            ET_LOG(
                Error,
                "Option %s is deprecated and no longer orders methods across "
                "threads. See the comment at this check for what to do instead.",
                kUseSharedCudaStream);
            return Error::NotSupported;
          }
        } else {
          ET_LOG(Error, "Option %s must be a boolean.", kUseSharedCudaStream);
          return Error::InvalidArgument;
        }
      } else if (std::strcmp(option.key, kWeightSharingAcrossMethods) == 0) {
        if (auto* val = std::get_if<bool>(&option.value)) {
          set_weight_sharing_across_methods(*val);
        } else {
          ET_LOG(
              Error,
              "Option %s must be a boolean.",
              kWeightSharingAcrossMethods);
          return Error::InvalidArgument;
        }
      } else if (std::strcmp(option.key, kEnableCudaGraphForMethod) == 0) {
#if defined(EXECUTORCH_USE_HIP)
        // HIP ignores the graph-instantiation flag required by this path.
        ET_LOG(
            Error,
            "Option %s is not supported on ROCm: HIP ignores graph "
            "instantiation flags.",
            kEnableCudaGraphForMethod);
        return Error::NotSupported;
#else
        if (auto* val = std::get_if<std::array<char, kMaxOptionValueLength>>(
                &option.value)) {
          set_cuda_graph_method(*val);
        } else {
          ET_LOG(
              Error,
              "Option %s must be a method name string.",
              kEnableCudaGraphForMethod);
          return Error::InvalidArgument;
        }
#endif
      }
    }
    return Error::Ok;
  }

  Error get_option(
      ET_UNUSED BackendOptionContext& context,
      ET_UNUSED executorch::runtime::Span<BackendOption>& backend_options)
      override {
    return Error::Ok;
  }

  // Once per loaded binary blob
  Result<DelegateHandle*> init(
      BackendInitContext& context,
      FreeableBuffer* processed, // This will be a empty buffer
      ArrayRef<CompileSpec> compile_specs // This will be my empty list
  ) const override {
    std::string method_name;
    std::optional<OffGraphKVStepWidth> kv_step_width;
    for (const CompileSpec& spec : compile_specs) {
      if (std::strcmp(spec.key, kOffGraphKVStepWidthSpec) == 0) {
        auto parsed = parse_offgraph_kv_step_width(std::string_view(
            static_cast<const char*>(spec.value.buffer), spec.value.nbytes));
        ET_CHECK_OK_OR_RETURN_ERROR(parsed.error());
        kv_step_width = parsed.get();
      }
      if (std::strcmp(spec.key, "method_name") == 0) {
        method_name.assign(
            static_cast<const char*>(spec.value.buffer),
            spec.value.nbytes); // no nullptr guarantee, so pass size
      }
    }

    std::string so_blob_key;
    std::string weights_blob_key;
    CudaWeightCache::Metadata fqn_weights;
    const bool has_fqn_weights =
        CudaWeightCache::is_serialized(processed->data(), processed->size());
    if (has_fqn_weights) {
      ET_CHECK_OK_OR_RETURN_ERROR(
          CudaWeightCache::parse(
              processed->data(), processed->size(), fqn_weights),
          "Malformed CUDA FQN weight metadata");
      size_t variant_index = 0;
      bool uses_ptx_fallback = false;
      uint32_t current_sm = 0;
      if (!(fqn_weights.variants.size() == 1 &&
            fqn_weights.variants[0].target_sm == 0)) {
#if defined(EXECUTORCH_USE_HIP)
        ET_LOG(
            Error,
            "Multi-SM CUDA AOTI metadata is not supported by the ROCm runtime");
        return Error::NotSupported;
#else
        int device_index = 0;
        cudaDeviceProp device_properties{};
        ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&device_index));
        ET_CUDA_CHECK_OR_RETURN_ERROR(
            cudaGetDeviceProperties(&device_properties, device_index));
        current_sm = static_cast<uint32_t>(
            device_properties.major * 10 + device_properties.minor);
        ET_CHECK_OK_OR_RETURN_ERROR(
            CudaWeightCache::select_variant(
                fqn_weights, current_sm, variant_index, uses_ptx_fallback),
            "Failed to select a CUDA AOTI variant for sm%u",
            current_sm);
#endif
      }
      const auto& variant = fqn_weights.variants[variant_index];
      so_blob_key = variant.so_blob_key;
      if (variant.target_sm == 0) {
        ET_LOG(Info, "Selected untargeted CUDA AOTI variant");
      } else if (uses_ptx_fallback) {
        ET_LOG(
            Info,
            "Selected sm%u CUDA AOTI PTX fallback (compute_%u) for sm%u",
            variant.target_sm,
            variant.ptx_compute,
            current_sm);
      } else {
        ET_LOG(
            Info, "Selected native sm%u CUDA AOTI variant", variant.target_sm);
      }
    } else {
      ET_CHECK_OK_OR_RETURN_ERROR(
          executorch::backends::aoti::resolve_blob_keys(
              processed, method_name, so_blob_key, weights_blob_key),
          "Malformed named-data key payload");
    }

    const NamedDataMap* named_data_map = context.get_named_data_map();
    auto aoti_dso_buffer = named_data_map->get_data(so_blob_key.c_str());
    ET_CHECK_OR_RETURN_ERROR(
        aoti_dso_buffer.ok(),
        Internal,
        "Failed to get data for key %s: 0x%x",
        so_blob_key.c_str(),
        static_cast<uint32_t>(aoti_dso_buffer.error()));

    // Generate a unique temporary file path. so_blob_key already selected the
    // blob above and is deliberately kept out of the filename: it can be an
    // untrusted, variable-length payload key, so embedding it risks path
    // traversal and cross-key collisions. Uniqueness comes from the pid
    // (across processes) and an atomic counter (within a process, since two
    // identical CUDA partitions would otherwise clobber each other's .so).
    static std::atomic<uint64_t> so_file_counter{0};
    filesystem::path temp_dir = filesystem::temp_directory_path();
    filesystem::path so_path = temp_dir /
        ("executorch_cuda_" + to_string(get_process_id()) + "_" +
         to_string(so_file_counter.fetch_add(1, std::memory_order_relaxed)) +
         ".so");

    // Create a temporary file
    ofstream outfile(so_path, ios::binary);

    // Write the ELF buffer to the temporary file
    ET_LOG(
        Info,
        "Writing %zu bytes to %s",
        aoti_dso_buffer->size(),
        so_path.c_str());

    outfile.write(
        static_cast<const char*>(aoti_dso_buffer->data()),
        aoti_dso_buffer->size());

    ET_CHECK_OR_RETURN_ERROR(
        outfile, AccessFailed, "Failed to write to file %s", so_path.c_str());

    // Finish writing the file to disk
    outfile.close();

    // Free the buffer immediately after writing to disk
    aoti_dso_buffer->Free();
    // Load the lib
    Result<void*> lib_handle_res = load_library(so_path);
    if (!lib_handle_res.ok()) {
      return lib_handle_res.error();
    }
    void* lib_handle = lib_handle_res.get();

    processed->Free();

    // Create handle and load function pointers into it
    cuda::CudaDelegateHandle* handle = new cuda::CudaDelegateHandle();
    handle->so_handle = lib_handle;
    handle->so_path = so_path.string();
    handle->method_name = method_name;

    // Load function pointers specific to this handle's shared library
    ET_CHECK_OK_OR_RETURN_ERROR(
        load_function_pointers_into_handle(lib_handle, handle));

    AOTInductorModelContainerHandle container_handle = nullptr;

    ET_CHECK_OK_OR_RETURN_ERROR(
        handle->create_with_device(&container_handle, 1, "cuda", nullptr));

    ET_LOG(Info, "container_handle = %p", container_handle);

    handle->container_handle = container_handle;

    // Runtime-owned off-graph buffers must capture their AOTI names before
    // the serialized constants update installs the ordinary weight set.
    //
    // The cache is found by the key the runner published as a runtime spec.
    // No key means an in-graph model, whose KV state is ordinary mutable
    // buffers -- that is the whole of the off-graph/in-graph detection, and it
    // needs nothing from the user.
    //
    // The key alone does not decide it, though: a program whose storage was
    // lowered off-graph carries no bytes for it, so running one without a cache
    // would hand its kernels unbound memory. That requirement is read from the
    // program's own constants and checked whether or not a key was passed.
    const char* cache_key = nullptr;
    if (auto spec = context.get_runtime_spec<const char*>(
            ::executorch::extension::llm::cache::kCacheKeyOption);
        spec.ok() && spec.get() != nullptr && *spec.get() != '\0') {
      cache_key = spec.get();
    }
    const auto needs_offgraph_kv = requires_offgraph_kv_storage(*handle);
    ET_CHECK_OK_OR_RETURN_ERROR(needs_offgraph_kv.error());
    if (cache_key != nullptr) {
#if defined(EXECUTORCH_BUILD_EXTENSION_LLM)
      ET_CHECK_OK_OR_RETURN_ERROR(
          attach_offgraph_kv_cache(*handle, cache_key, kv_step_width));
#else
      ET_LOG(
          Error,
          "init: an off-graph KV cache was supplied, but this CUDA backend was "
          "built without EXECUTORCH_BUILD_EXTENSION_LLM");
      return Error::NotSupported;
#endif
    }
    ET_CHECK_OR_RETURN_ERROR(
        !needs_offgraph_kv.get() || handle->kv_cache != nullptr,
        InvalidArgument,
        "init: method '%s' keeps its KV cache off-graph but no cache was "
        "installed under the '%s' backend option (a CUDA backend built without "
        "EXECUTORCH_BUILD_EXTENSION_LLM cannot run it)",
        method_name.c_str(),
        ::executorch::extension::llm::cache::kCacheKeyOption);

    // Versioned artifacts load each (device, FQN) through the same process-wide
    // cross-method cache model used by the legacy path. The payload only adds
    // the tensor metadata needed to reconstruct independently named PTD blobs.
    if (has_fqn_weights) {
      ET_CHECK_OK_OR_RETURN_ERROR(
          fqn_weight_cache_.load(handle, named_data_map, fqn_weights));
    } else if (is_weight_sharing_across_methods_enabled()) {
      ET_CHECK_OK_OR_RETURN_ERROR(load_constants_with_cache(
          handle, named_data_map, method_name, weights_blob_key));
    } else {
      ET_CHECK_OK_OR_RETURN_ERROR(
          load_constants_legacy(handle, named_data_map, weights_blob_key));
    }

    // Handles on one thread share that thread's stream, so one delegate's
    // output is ordered against the next one's read, which a stream per handle
    // left unordered. cudaStreamPerThread is a different stream on each host
    // thread, so this orders delegates called from the same thread and not
    // delegates called from different ones. The TensorRT delegate falls back to
    // the same stream, in its executorch backend, so a split program on one
    // thread is ordered too.
    handle->cuda_stream = cudaStreamPerThread;
    ET_LOG(
        Info,
        "Using the per-thread CUDA stream %p for method %s",
        handle->get_cuda_stream(),
        method_name.c_str());

    // Initialize CUDA graph state if enabled for this method.
    if (should_use_cuda_graph_for_method(method_name)) {
      handle->cuda_graph_state.phase = CudaGraphPhase::Warmup;
      handle->cuda_graph_state.warmup_remaining = kCudaGraphWarmupSteps;
      ET_LOG(
          Info,
          "CUDA graph enabled for method '%s' (warmup=%d)",
          method_name.c_str(),
          kCudaGraphWarmupSteps);
    }

    mutable_state_note_handle(handle);

    live_handles_.fetch_add(1, std::memory_order_acq_rel);

    return (DelegateHandle*)handle; // Return the handle post-processing
  }

  // Once per execution
  Error execute(
      BackendExecutionContext& context,
      DelegateHandle* handle_,
      Span<EValue*> args) const override {
    cuda::CudaDelegateHandle* handle = (cuda::CudaDelegateHandle*)handle_;

    size_t n_inputs;
    handle->get_num_inputs(handle->container_handle, &n_inputs);

    size_t n_outputs;
    handle->get_num_outputs(handle->container_handle, &n_outputs);

    // Run on the caller-selected stream when one is active on this thread (e.g.
    // a CUDA green-context stream), otherwise the per-thread stream. Every
    // kernel and boundary copy reads getCurrentCUDAStream, so installing the
    // choice here routes the whole execution; restore the prior selection on
    // return so a caller stream does not linger for later work on this thread.
    const std::optional<cudaStream_t> caller_stream =
        executorch::extension::cuda::getCallerStream();

    // Replaying a captured graph on a caller-provided stream is not itself a
    // CUDA error, but the static buffers this path pins are shared by every
    // replay, so two callers on two streams would race over them. Refused
    // rather than synchronized, which predates this change.
    ET_CHECK_OR_RETURN_ERROR(
        !(caller_stream &&
          handle->cuda_graph_state.phase != CudaGraphPhase::Disabled),
        NotSupported,
        "CUDA graph is not supported together with a caller-provided CUDA stream.");

    // Snapshot the prior selection without creating one (peek, not get), so the
    // restore is exact and we don't leak a stream just to snapshot.
    std::optional<cudaStream_t> prev_stream;
    if (caller_stream) {
      prev_stream = peekCurrentCUDAStream(0);
    }
    setCurrentCUDAStream(caller_stream.value_or(handle->get_cuda_stream()), 0);
    executorch::backends::aoti::ScopeGuard restore_stream([&]() noexcept {
      if (!caller_stream) {
        return;
      }
      if (prev_stream) {
        setCurrentCUDAStream(*prev_stream, 0);
      } else {
        clearCurrentCUDAStream(0);
      }
    });

    size_t n_io_sum = 0;
    ET_CHECK_OR_RETURN_ERROR(
        !c10::add_overflows(n_inputs, n_outputs, &n_io_sum) &&
            n_io_sum == args.size(),
        InvalidArgument,
        "number of user input %zd and output %zd generated from AOT Inductor does not match ET runner's %zd. Exit.",
        n_inputs,
        n_outputs,
        args.size())

    // Verify device info on all memory-planned, ET-driven IO tensors.
    // All input and output tensors should have device_type = CUDA, which
    // is set during serialization by PropagateDevicePass based on the
    // target_device compile spec from CudaPartitioner.
    //
    // Under device memory planning, these tensors are GPU-resident: their
    // storage lives in a planned CUDA arena. The backend wraps them in place
    // and performs no host<->device copies — graph-level et_copy ops handle
    // any host<->device transfers outside the delegate.
    for (size_t i = 0; i < n_inputs + n_outputs; i++) {
      auto* tensor = &(args[i]->toTensor());
      auto device_type = tensor->unsafeGetTensorImpl()->device_type();
      ET_CHECK_OR_RETURN_ERROR(
          device_type == executorch::runtime::etensor::DeviceType::CUDA,
          InvalidArgument,
          "Tensor %zu expected device_type=CUDA (1), got %d. "
          "Device info may not be properly propagated from CudaPartitioner.",
          i,
          static_cast<int>(device_type));

      // device_type above is only metadata. Also verify the storage actually
      // lives in CUDA device memory, so a CUDA-typed tensor that is secretly
      // backed by host memory is caught here instead of corrupting the run.
      const void* data_ptr = tensor->const_data_ptr();
      if (data_ptr != nullptr) {
        cudaPointerAttributes attributes{};
        const cudaError_t attr_err =
            cudaPointerGetAttributes(&attributes, data_ptr);
        ET_CHECK_OR_RETURN_ERROR(
            attr_err == cudaSuccess &&
                (attributes.type == cudaMemoryTypeDevice ||
                 attributes.type == cudaMemoryTypeManaged),
            InvalidArgument,
            "Tensor %zu has device_type=CUDA but its data pointer %p is not "
            "backed by CUDA device memory (cudaPointerGetAttributes err=%d, "
            "cudaMemoryType=%d).",
            i,
            data_ptr,
            static_cast<int>(attr_err),
            static_cast<int>(attributes.type));
      }
    }

    // The cache has to grow, and its storage be rebound, before the compiled
    // program reads it -- and nothing inside that program can call back out, so
    // the step is driven from here rather than by the runner.
    int64_t kv_step_width = 0;
    if (handle->kv_cache != nullptr) {
      const auto step_width = read_offgraph_kv_step_width(
          handle->kv_step_width, Span<EValue*>(args.data(), n_inputs));
      ET_CHECK_OK_OR_RETURN_ERROR(step_width.error());
      kv_step_width = step_width.get();
      // Growth is ordered on the stream this step runs on, which was just
      // installed as current above.
      const Result<cudaStream_t> step_stream = getCurrentCUDAStream(0);
      ET_CHECK_OK_OR_RETURN_ERROR(step_stream.error());
      ET_CHECK_OK_OR_RETURN_ERROR(
          handle->kv_cache->prepare_step(kv_step_width, step_stream.get()));
      ET_CHECK_OK_OR_RETURN_ERROR(handle->kv_cache->rebind_for_execute(handle));
    }
    ET_CHECK_OK_OR_RETURN_ERROR(mutable_state_rebind_for_execute(handle));

    // ---------------------------------------------------------------
    // CUDA graph REPLAY path — skip all tensor setup and just replay
    // ---------------------------------------------------------------
    if (handle->cuda_graph_state.phase == CudaGraphPhase::Replay) {
      Result<cudaStream_t> csr = getCurrentCUDAStream(0);
      ET_CHECK_OK_OR_RETURN_ERROR(csr.error());
      cudaStream_t cs = csr.get();

      // Copy new input data (GPU-resident) into static input buffers (D2D)
      for (size_t i = 0; i < n_inputs; i++) {
        auto* et_input = &(args[i]->toTensor());
        ET_CHECK_OR_RETURN_ERROR(
            et_input->nbytes() ==
                handle->cuda_graph_state.static_input_nbytes[i],
            InvalidArgument,
            "CUDA graph replay: input %zu size mismatch (expected %zu, got %zu)",
            i,
            handle->cuda_graph_state.static_input_nbytes[i],
            et_input->nbytes());
        ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync(
            handle->cuda_graph_state.static_input_ptrs[i],
            et_input->const_data_ptr(),
            handle->cuda_graph_state.static_input_nbytes[i],
            cudaMemcpyDeviceToDevice,
            cs));
      }

      // Replay the captured graph
      cudaError_t gerr =
          cudaGraphLaunch(handle->cuda_graph_state.graph_exec, cs);
      ET_CHECK_OR_RETURN_ERROR(
          gerr == cudaSuccess,
          Internal,
          "cudaGraphLaunch failed: %s",
          cudaGetErrorString(gerr));

      // Copy outputs from static buffers into the planned GPU ET buffers (D2D)
      for (size_t i = 0; i < n_outputs; i++) {
        auto* et_output = &(args[i + n_inputs]->toTensor());
        ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync(
            et_output->mutable_data_ptr(),
            handle->cuda_graph_state.static_output_ptrs[i],
            handle->cuda_graph_state.static_output_nbytes[i],
            cudaMemcpyDeviceToDevice,
            cs));
      }
      ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(cs));

      return commit_offgraph_kv_step(handle, kv_step_width);
    }

    // ---------------------------------------------------------------
    // Normal path (also used for WARMUP and CAPTURE phases)
    // ---------------------------------------------------------------
    bool is_capture_step =
        (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup &&
         handle->cuda_graph_state.warmup_remaining == 0);

    // Wrrap them as SlimTensors for AOTI execution without any host<->device
    // copies.
    std::vector<SlimTensor*> slim_inputs(n_inputs);
    std::vector<SlimTensor*> slim_outputs(n_outputs);

    // Undoes a capture attempt that an early return would otherwise abandon.
    //
    // The buffers this attempt pinned would otherwise stay in their vectors
    // with the phase still at warmup and no steps left, so the next call
    // captures again and appends a second set. Replay then reads the second set
    // while the input copies target the first, and every execute returns
    // whatever those buffers held at capture time, with nothing reporting an
    // error.
    //
    // Ending the capture itself is a separate guard, declared after the tensor
    // cleanup below so that it runs before it: freeing a device buffer on a
    // still-capturing stream fails with invalid argument and leaks the block.
    //
    // Disarmed once the capture step has fully succeeded.
    class CaptureGuard {
     public:
      ~CaptureGuard() {
        if (state_ == nullptr) {
          return;
        }
        // Free what this attempt pinned and disable graphs for this method, so
        // a capture that cannot succeed costs one error rather than one on
        // every fourth call for the life of the process. Eager execution is
        // correct, just slower.
        for (void* ptr : state_->static_input_ptrs) {
          (void)cudaFree(ptr);
        }
        state_->static_input_ptrs.clear();
        state_->static_output_ptrs.clear();
        state_->static_input_nbytes.clear();
        state_->static_output_nbytes.clear();
        // Same order as ~CudaGraphState: the exec depends on the graph.
        if (state_->graph_exec != nullptr) {
          (void)cudaGraphExecDestroy(state_->graph_exec);
          state_->graph_exec = nullptr;
        }
        if (state_->graph != nullptr) {
          (void)cudaGraphDestroy(state_->graph);
          state_->graph = nullptr;
        }
        state_->phase = CudaGraphPhase::Disabled;
        state_->warmup_remaining = 0;
        (void)cudaGetLastError();
      }
      // Before capture begins. From here a failure still unwinds the pinned
      // buffers.
      void arm(cuda::CudaGraphState* state) {
        state_ = state;
      }
      void disarm() {
        state_ = nullptr;
      }

     private:
      cuda::CudaGraphState* state_ = nullptr;
    } capture_guard;

    if (is_capture_step) {
      capture_guard.arm(&handle->cuda_graph_state);
    }

    // Process input tensors: wrap the GPU-resident ETensor buffers directly.
    for (size_t i = 0; i < n_inputs; i++) {
      auto* et_input = &(args[i]->toTensor());

      // CAPTURE step: allocate persistent static GPU buffers and seed them
      // from the GPU-resident ET inputs (D2D).
      if (is_capture_step) {
        size_t nbytes = et_input->nbytes();

        void* static_ptr = nullptr;
        cudaError_t merr = cudaMalloc(&static_ptr, nbytes);
        ET_CHECK_OR_RETURN_ERROR(
            merr == cudaSuccess,
            Internal,
            "cudaMalloc for static input %zu failed: %s",
            i,
            cudaGetErrorString(merr));

        // Tracked before the seeding copy, so a failed copy still unwinds
        // through the guard instead of leaking this allocation.
        handle->cuda_graph_state.static_input_ptrs.push_back(static_ptr);
        handle->cuda_graph_state.static_input_nbytes.push_back(nbytes);

        ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpy(
            static_ptr,
            et_input->const_data_ptr(),
            nbytes,
            cudaMemcpyDeviceToDevice));

        slim_inputs[i] = make_slimtensor_from_blob_with_etensor_metadata(
            static_ptr, et_input);
        continue;
      }

      // Normal path: wrap the GPU buffer in place (zero-copy, non-owning).
      slim_inputs[i] = make_slimtensor_from_blob_with_etensor_metadata(
          const_cast<void*>(et_input->const_data_ptr()), et_input);
    }

    // Process output tensors: wrap the GPU-resident ET output buffers as
    // non-owning SlimTensors so AOTI can write into the planned slot directly
    // when it does not allocate its own output. Save pre-run handles to detect
    // when AOTI replaces them with its own allocation.
    std::vector<SlimTensor*> pre_run_outputs(n_outputs, nullptr);
    for (size_t i = 0; i < n_outputs; i++) {
      auto* et_output = &(args[i + n_inputs]->toTensor());
      slim_outputs[i] = make_slimtensor_from_blob_with_etensor_metadata(
          const_cast<void*>(et_output->const_data_ptr()), et_output);
      pre_run_outputs[i] = slim_outputs[i];
    }

    bool run_called = false;

    // Scope guard: deletes any non-null slim_outputs on exit. Normal paths
    // null entries as they take ownership, so the guard only fires on
    // early-return error paths. Also cleans up inputs if run() was never
    // called (run() steals them via internal RAII).
    executorch::backends::aoti::ScopeGuard cleanup([&]() noexcept {
      if (!run_called) {
        delete_slimtensor_vector(slim_inputs);
      }
      for (size_t i = 0; i < slim_outputs.size(); i++) {
        if (slim_outputs[i]) {
          delete slim_outputs[i];
        }
      }
    });

    // Ends a capture that an early return would otherwise leave running.
    // Declared after `cleanup` so it is destroyed before it: a device buffer
    // freed on a still-capturing stream fails with invalid argument and leaks
    // the block.
    //
    // Leaving the stream capturing matters because handles share the per-thread
    // stream: the next delegate on this thread would have its kernels captured
    // instead of run, and later synchronizes would fail.
    class EndCaptureGuard {
     public:
      ~EndCaptureGuard() {
        if (stream_ == nullptr) {
          return;
        }
        cudaGraph_t abandoned = nullptr;
        if (cudaStreamEndCapture(stream_, &abandoned) == cudaSuccess) {
          if (abandoned != nullptr) {
            (void)cudaGraphDestroy(abandoned);
          }
        } else {
          // Clears the sticky error so the next unrelated CUDA call on this
          // thread does not inherit it. This clears whatever error is pending,
          // not only the one from above.
          (void)cudaGetLastError();
        }
      }
      void arm(cudaStream_t stream) {
        stream_ = stream;
      }
      void disarm() {
        stream_ = nullptr;
      }

     private:
      cudaStream_t stream_ = nullptr;
    } end_capture_guard;

    // Run the AOTI container.
    // NOTE: run() steals input handles (RAII wraps them at the start of
    // run_impl) and may replace output handles with its own.
    Result<cudaStream_t> cuda_stream_ret = getCurrentCUDAStream(0);
    ET_CHECK_OK_OR_RETURN_ERROR(cuda_stream_ret.error());
    cudaStream_t cuda_stream = cuda_stream_ret.get();

    if (is_capture_step) {
      // ----- CUDA graph CAPTURE -----
      ET_LOG(
          Info,
          "CUDA graph: beginning stream capture for '%s'",
          handle->method_name.c_str());

      cudaError_t cerr =
          cudaStreamBeginCapture(cuda_stream, cudaStreamCaptureModeRelaxed);
      ET_CHECK_OR_RETURN_ERROR(
          cerr == cudaSuccess,
          Internal,
          "cudaStreamBeginCapture failed: %s",
          cudaGetErrorString(cerr));
      end_capture_guard.arm(cuda_stream);
    }

    AOTIRuntimeError error = handle->run(
        handle->container_handle,
        reinterpret_cast<Tensor**>(slim_inputs.data()),
        n_inputs,
        reinterpret_cast<Tensor**>(slim_outputs.data()),
        n_outputs,
        static_cast<void*>(cuda_stream),
        nullptr);
    run_called = true;

    // Delete orphaned pre-created outputs that run() replaced.
    // Must happen before the error check — if run() fails after
    // replacing some outputs, the originals would otherwise leak.
    for (size_t i = 0; i < n_outputs; i++) {
      if (pre_run_outputs[i] != slim_outputs[i]) {
        delete pre_run_outputs[i];
      }
    }

    ET_CHECK_OR_RETURN_ERROR(
        error == Error::Ok,
        Internal,
        "AOTInductorModelContainerRun failed with error code %d",
        error);

    if (is_capture_step) {
      // End capture → instantiate graph
      cudaError_t gerr =
          cudaStreamEndCapture(cuda_stream, &handle->cuda_graph_state.graph);
      // The stream has left capture either way, so the guard must not end it
      // again; the state guard below still unwinds what the attempt pinned.
      end_capture_guard.disarm();
      ET_CHECK_OR_RETURN_ERROR(
          gerr == cudaSuccess,
          Internal,
          "cudaStreamEndCapture failed: %s",
          cudaGetErrorString(gerr));

      gerr = cudaGraphInstantiate(
          &handle->cuda_graph_state.graph_exec,
          handle->cuda_graph_state.graph,
          cudaGraphInstantiateFlagAutoFreeOnLaunch);
      ET_CHECK_OR_RETURN_ERROR(
          gerr == cudaSuccess,
          Internal,
          "cudaGraphInstantiate failed: %s",
          cudaGetErrorString(gerr));

      // Record static output pointers (stable under graph replay). Releasing
      // them from slim_outputs here, before the copies below, keeps the cleanup
      // guard from deleting buffers the AOTI runtime owns if one of those
      // copies fails.
      for (size_t i = 0; i < n_outputs; i++) {
        SlimTensor* out = slim_outputs[i];
        handle->cuda_graph_state.static_output_ptrs.push_back(out->data_ptr());
        handle->cuda_graph_state.static_output_nbytes.push_back(out->nbytes());
        slim_outputs[i] = nullptr;
      }

      handle->cuda_graph_state.phase = CudaGraphPhase::Replay;
      ET_LOG(
          Info,
          "CUDA graph: captured and instantiated for '%s'",
          handle->method_name.c_str());

      // Replay once to actually produce output (capture doesn't execute)
      gerr = cudaGraphLaunch(handle->cuda_graph_state.graph_exec, cuda_stream);
      ET_CHECK_OR_RETURN_ERROR(
          gerr == cudaSuccess,
          Internal,
          "cudaGraphLaunch (first replay) failed: %s",
          cudaGetErrorString(gerr));

      // Copy capture-step outputs into the planned GPU ET buffers (D2D).
      for (size_t i = 0; i < n_outputs; i++) {
        auto* et_output = &(args[i + n_inputs]->toTensor());
        ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync(
            et_output->mutable_data_ptr(),
            handle->cuda_graph_state.static_output_ptrs[i],
            handle->cuda_graph_state.static_output_nbytes[i],
            cudaMemcpyDeviceToDevice,
            cuda_stream));
      }
      ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(cuda_stream));

      // Last failure point is behind us, so the captured state is now the state
      // the next call should replay from.
      capture_guard.disarm();
      return commit_offgraph_kv_step(handle, kv_step_width);
    }

    // ----- Normal / WARMUP execution continues here -----

    // Decrement warmup counter if in warmup phase
    if (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup &&
        handle->cuda_graph_state.warmup_remaining > 0) {
      handle->cuda_graph_state.warmup_remaining--;
      ET_LOG(
          Info,
          "CUDA graph warmup: %d steps remaining for '%s'",
          handle->cuda_graph_state.warmup_remaining,
          handle->method_name.c_str());
    }

    // Land each output into its planned GPU ET buffer.
    //
    // Before run(), each slim_outputs[i] was a non-owning view over the
    // planned GPU ET output buffer, and we recorded that pointer in
    // pre_run_outputs[i]. AOTInductorModelContainerRun may either:
    //   (a) write directly into the buffer we passed (leaving our handle in
    //       place), or
    //   (b) allocate its own output in its caching allocator and overwrite
    //       slim_outputs[i] with that new handle.
    // Comparing the post-run handle against the recorded pre-run handle tells
    // us which happened.
    for (size_t i = 0; i < n_outputs; i++) {
      auto* et_output = &(args[i + n_inputs]->toTensor());
      if (pre_run_outputs[i] == slim_outputs[i]) {
        // Case (a): AOTI wrote directly into our planned ET buffer. The result
        // is already in place — nothing to copy, just drop the view wrapper.
        delete slim_outputs[i];
      } else {
        // Case (b): AOTI returned its own buffer. D2D copy the result into the
        // planned GPU ET buffer, then free AOTI's buffer.
        ET_CHECK_OK_OR_RETURN_ERROR(
            copy_slimtensor_to_device_etensor_async(
                slim_outputs[i], et_output, cuda_stream),
            "Failed to D2D copy GPU output %zu into ETensor",
            i);
        delete slim_outputs[i];
      }
      slim_outputs[i] = nullptr;
    }

    return commit_offgraph_kv_step(handle, kv_step_width);
  }

  void destroy(DelegateHandle* handle_) const override {
    if (handle_ == nullptr) {
      return;
    }
    cuda::CudaDelegateHandle* handle = (cuda::CudaDelegateHandle*)handle_;

    mutable_state_forget_handle(handle);
    if (handle->kv_cache != nullptr) {
      handle->kv_cache->forget_handle(handle);
    }

    // NOTE: AOTInductorModelContainerDelete does not work correctly with
    // multiple .so files. Deleting one container frees shared resources,
    // which causes segmentation faults when attempting to delete other
    // containers. As a workaround, we skip explicit container deletion and
    // defer cleanup to the OS. The corresponding shared library must remain
    // loaded as well: the leaked container still owns objects whose code and
    // process-wide state live in that library, so dlclose/FreeLibrary can
    // invalidate them and crash during multi-method teardown.
    // TODO(gasoonjia): Find a proper solution for safe container deletion.
    // AOTInductorModelContainerDelete(handle->container_handle);

    // Remove the temporary shared library file
    if (!handle->so_path.empty()) {
      std::error_code remove_error;
      std::filesystem::remove(handle->so_path, remove_error);
      ET_CHECK_OR_LOG_ERROR(
          !remove_error,
          "Failed to remove temporary shared library %s: %s",
          handle->so_path.c_str(),
          remove_error.message().c_str());
    }

    delete handle;

    // The allocator lets the device pool keep freed memory so that repeated
    // delegate execution does not pay to map it again. Nothing is running on
    // this backend once the last handle is gone, so hand that memory back
    // rather than hold it for the life of the process.
    if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
      // Only frees the driver has already observed can be released, so without
      // this the trim gives back nothing rather than less. A device wait rather
      // than a stream wait because the frees went to whichever stream was
      // current when the storage was released, which is not necessarily one
      // this handle recorded, and teardown is not bound to the thread that ran
      // the method either.
      const cudaError_t sync_err = cudaDeviceSynchronize();
      if (sync_err != cudaSuccess) {
        ET_LOG(
            Error,
            "cudaDeviceSynchronize before releasing the pool failed: %s.",
            cudaGetErrorString(sync_err));
        (void)cudaGetLastError();
      }
      CudaAllocator::release_cached_memory(-1);
    }
  }

 private:
  mutable std::mutex cuda_graph_method_mutex_;
  std::string cuda_graph_method_;

  // Whether to enable cross-method caching for legacy dense-blob artifacts.
  // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default
  // OFF; versioned FQN artifacts do not consult this option.
  std::atomic<bool> weight_sharing_across_methods_{false};

  // Delegates alive right now. The device memory pool is shared, so it can only
  // be released once none of them are left.
  mutable std::atomic<size_t> live_handles_{0};

  // ---------------------------------------------------------------
  // Per-weight constant cache.
  //
  // Maintains a singleton FQN → AtenTensorHandle cache across methods.
  // When loading constants for a method, constants already in the cache
  // are reused (zero-copy via update_user_managed_constant_buffer_pairs).
  // Only constants not in the cache are loaded from the blob and added
  // to the cache. This avoids duplicate GPU allocations when multiple
  // methods (e.g., prefill/decode) share the same weights.
  //
  // ASSUMPTIONS / LIMITATIONS:
  //   * Constants with the same FQN across methods are assumed to be the
  //     SAME logical tensor (i.e. the same parameter/buffer of the same
  //     source model). We validate shape/dtype/strides/device on every
  //     reuse to catch silent mismatches (see check_cached_constant_match
  //     below). However, we cannot detect two unrelated models that
  //     happen to share an FQN.
  //   * Constants are assumed to be IMMUTABLE (parameters or read-only
  //     buffers). The AOTI shim today does not expose a mutability bit
  //     through GetConstantOriginalFQN, so we cannot detect or refuse
  //     to share mutable buffers (for example, runtime caches). If a
  //     model exports the same FQN as a mutable buffer in multiple
  //     methods, mutations from one method WILL be visible to the other
  //     through the shared GPU memory. Callers that need isolated mutable
  //     state for shared FQNs must opt into cuda_mutable_state or use
  //     distinct FQNs.
  //     TODO: when AOTInductor exposes a constant-type / mutability
  //     query, refuse to share entries that are not PARAMETER or
  //     non-mutable BUFFER.
  // ---------------------------------------------------------------

  // Validates that a cached constant tensor is compatible with what the
  // new container expects for the same FQN (i.e. same dtype, dim, sizes,
  // strides, and device). Both handles point to SlimTensors in our shim
  // layer, so we can introspect them directly.
  //
  // Returns Error::Ok on a match. On mismatch, logs the offending field
  // and returns Error::Internal so callers can chain via
  // ET_CHECK_OK_OR_RETURN_ERROR and fail loudly instead of silently
  // pointing the new container at a wrong-shape buffer.
  static Error check_cached_constant_match(
      const std::string& fqn,
      AtenTensorHandle cached_handle,
      AtenTensorHandle new_handle) {
    ET_CHECK_OR_RETURN_ERROR(
        cached_handle != nullptr && new_handle != nullptr,
        Internal,
        "Constant '%s': null AtenTensorHandle (cached=%p, new=%p)",
        fqn.c_str(),
        cached_handle,
        new_handle);

    auto* cached = reinterpret_cast<SlimTensor*>(cached_handle);
    auto* fresh = reinterpret_cast<SlimTensor*>(new_handle);

    ET_CHECK_OR_RETURN_ERROR(
        cached->dtype() == fresh->dtype(),
        Internal,
        "Constant '%s': dtype mismatch (cached=%d, new=%d)",
        fqn.c_str(),
        static_cast<int>(cached->dtype()),
        static_cast<int>(fresh->dtype()));

    ET_CHECK_OR_RETURN_ERROR(
        cached->dim() == fresh->dim(),
        Internal,
        "Constant '%s': dim mismatch (cached=%zu, new=%zu)",
        fqn.c_str(),
        cached->dim(),
        fresh->dim());

    auto cached_sizes = cached->sizes();
    auto fresh_sizes = fresh->sizes();
    for (size_t i = 0; i < cached->dim(); ++i) {
      ET_CHECK_OR_RETURN_ERROR(
          cached_sizes[i] == fresh_sizes[i],
          Internal,
          "Constant '%s': size mismatch at dim %zu (cached=%lld, new=%lld)",
          fqn.c_str(),
          i,
          static_cast<long long>(cached_sizes[i]),
          static_cast<long long>(fresh_sizes[i]));
    }
    auto cached_strides = cached->strides();
    auto fresh_strides = fresh->strides();
    for (size_t i = 0; i < cached->dim(); ++i) {
      ET_CHECK_OR_RETURN_ERROR(
          cached_strides[i] == fresh_strides[i],
          Internal,
          "Constant '%s': stride mismatch at dim %zu (cached=%lld, new=%lld)",
          fqn.c_str(),
          i,
          static_cast<long long>(cached_strides[i]),
          static_cast<long long>(fresh_strides[i]));
    }
    ET_CHECK_OR_RETURN_ERROR(
        cached->device().type() == fresh->device().type() &&
            cached->device().index() == fresh->device().index(),
        Internal,
        "Constant '%s': device mismatch (cached=%d:%d, new=%d:%d)",
        fqn.c_str(),
        static_cast<int>(cached->device().type()),
        cached->device().index(),
        static_cast<int>(fresh->device().type()),
        fresh->device().index());

    return Error::Ok;
  }

  // Load constants for a method using per-weight caching.
  // Returns Error::Ok on success.
  //
  // Flow:
  //   1. Enumerate this method's constants and their FQNs.
  //   2. For each constant:
  //      - If FQN is in shared_constant_tensors_ → reuse (cache hit).
  //      - Otherwise → mark as needing loading (cache miss).
  //   3. If all constants are cached → skip blob loading entirely.
  //      Otherwise → call update_constants_from_blob to load all, then
  //      extract and cache the new constants.
  //   4. For cached constants, call update_user_managed_constant_buffer_pairs
  //      to point the container to the shared GPU tensors.
  Error load_constants_with_cache(
      cuda::CudaDelegateHandle* handle,
      const NamedDataMap* named_data_map,
      const std::string& method_name,
      const std::string& weights_blob_key) const {
    // Check if the required APIs are available
    if (!handle->get_num_constants || !handle->get_constant_name ||
        !handle->get_constant_original_fqn || !handle->extract_constants_map ||
        !handle->update_user_managed_constant_buffer_pairs) {
      // Fall back to the legacy path
      return load_constants_legacy(handle, named_data_map, weights_blob_key);
    }

    // Step 1: Enumerate constants and partition into cached/uncached
    size_t num_constants = 0;
    ET_CHECK_OK_OR_RETURN_ERROR(
        handle->get_num_constants(handle->container_handle, &num_constants),
        "Failed to enumerate CUDA AOTI constants for method '%s'",
        method_name.c_str());
    if (num_constants == 0) {
      ET_LOG(Info, "No constants for method '%s'", method_name.c_str());
      return Error::Ok;
    }

    // Build FQN → internal_name mapping and determine cache hits/misses.
    std::unordered_map<std::string, std::string> fqn_to_name;
    std::vector<std::string> uncached_fqns;

    // Phase 1 (lock-free): enumerate constants from the container.
    for (size_t i = 0; i < num_constants; i++) {
      const char* name = nullptr;
      const char* fqn = nullptr;
      handle->get_constant_name(handle->container_handle, i, &name);
      handle->get_constant_original_fqn(handle->container_handle, i, &fqn);
      if (name && fqn && fqn[0] != '\0') {
        fqn_to_name[fqn] = name;
      }
    }

    // Names this method must load itself because the cached tensor under the
    // same name belongs to a different constant.
    std::unordered_set<std::string> not_shared;

    // Phase 2 (locked): pure cache lookup against shared_constant_tensors_.
    // A cache hit here is provisional: the name matches, but whether the
    // tensors match is only known after extraction, so a hit can still be
    // rejected below. Only the path that loads the blob can check this. When
    // every name is already cached there is nothing to compare against
    // without re-uploading the blob, which costs too much device memory to
    // do on every method, so that case still shares on the name alone.
    {
      std::lock_guard<std::mutex> guard(shared_constants_mutex_);
      for (const auto& [fqn, _] : fqn_to_name) {
        if (shared_constant_tensors_.find(fqn) ==
            shared_constant_tensors_.end()) {
          uncached_fqns.push_back(fqn);
        }
      }
    }

    size_t num_cached = fqn_to_name.size() - uncached_fqns.size();
    ET_LOG(
        Info,
        "Method '%s': %zu constants, %zu cached, %zu uncached",
        method_name.c_str(),
        fqn_to_name.size(),
        num_cached,
        uncached_fqns.size());

    // Step 2: Load uncached constants from blob (if any).
    std::unordered_map<std::string, AtenTensorHandle> extracted_map;

    if (!uncached_fqns.empty()) {
      // Need to load from blob — use update_constants_from_blob for all,
      // then extract the new constants into the cache.
      auto buffer_res = named_data_map->get_data(weights_blob_key.c_str());

      ET_CHECK_OR_RETURN_ERROR(
          buffer_res.ok() && handle->update_constants_from_blob != nullptr,
          NotFound,
          "weights_blob '%s' not found or update fn is null",
          weights_blob_key.c_str());

      ET_LOG(
          Info,
          "Loading constants from blob '%s' for method '%s'",
          weights_blob_key.c_str(),
          method_name.c_str());
      const void* weights_blob = buffer_res->data();
      ET_CHECK_OK_OR_RETURN_ERROR(
          handle->update_constants_from_blob(
              handle->container_handle,
              static_cast<const uint8_t*>(weights_blob)),
          "update_constants_from_blob failed for method '%s'",
          method_name.c_str());
      ET_CUDA_CHECK_OR_RETURN_ERROR(cudaDeviceSynchronize());
      buffer_res->Free();

      // Extract all constants from the freshly-loaded container.
      ET_CHECK_OK_OR_RETURN_ERROR(
          handle->extract_constants_map(
              handle->container_handle,
              reinterpret_cast<AOTInductorConstantMapHandle>(&extracted_map),
              /*use_inactive=*/false),
          "Failed to extract constants from '%s'",
          method_name.c_str());

      // Validate cache hits against the freshly-extracted tensors, and
      // populate the cache with newly-loaded entries.
      {
        std::lock_guard<std::mutex> guard(shared_constants_mutex_);
        for (const auto& [fqn, _] : fqn_to_name) {
          auto extracted_it = extracted_map.find(fqn);
          if (extracted_it == extracted_map.end()) {
            // Container did not surface this FQN — skip; the user-managed
            // pair build below will simply omit it.
            continue;
          }
          auto cached_it = shared_constant_tensors_.find(fqn);
          if (cached_it == shared_constant_tensors_.end()) {
            // New constant — add to cache.
            shared_constant_tensors_[fqn] = extracted_it->second;
          } else if (
              check_cached_constant_match(
                  fqn, cached_it->second, extracted_it->second) != Error::Ok) {
            // Same name, different tensor. AOTInductor names a lifted constant
            // by position within one compiled module, so two methods each own a
            // "_tensor_constant2" holding unrelated data, and a name is only
            // evidence of sharing when the tensors agree. Sharing this pair
            // would alias mismatched storage, so leave this method's own
            // constant in place and keep the cached one for whoever matches it.
            ET_LOG(
                Info,
                "Constant '%s' in method '%s' differs from the cached copy; "
                "using this method's own copy instead of sharing",
                fqn.c_str(),
                method_name.c_str());
            not_shared.insert(fqn);
          }
        }
        ET_LOG(
            Info,
            "Cached %zu new constants from method '%s' (total cache: %zu)",
            uncached_fqns.size(),
            method_name.c_str(),
            shared_constant_tensors_.size());
      }
    } else {
      // All constants are cached — skip blob loading entirely.
      // NOTE: in this branch we cannot independently verify the cache
      // against the new container's expectations (no extract source).
      // We rely on update_user_managed_constant_buffer_pairs below,
      // which the AOTI runtime validates internally.
      ET_LOG(
          Info,
          "All %zu constants cached — skipping blob load for method '%s'",
          fqn_to_name.size(),
          method_name.c_str());
    }

    // Step 3: Point the container to cached tensors via user_managed pairs
    if (num_cached > 0 || uncached_fqns.empty()) {
      std::vector<AOTInductorConstantMapEntry> pairs;
      {
        std::lock_guard<std::mutex> guard(shared_constants_mutex_);
        for (const auto& [fqn, internal_name] : fqn_to_name) {
          if (not_shared.count(fqn) != 0) {
            continue;
          }
          auto it = shared_constant_tensors_.find(fqn);
          if (it != shared_constant_tensors_.end()) {
            pairs.push_back({internal_name.c_str(), it->second});
          }
        }
      }

      if (!pairs.empty()) {
        ET_CHECK_OK_OR_RETURN_ERROR(
            handle->update_user_managed_constant_buffer_pairs(
                handle->container_handle,
                pairs.data(),
                pairs.size(),
                /*use_inactive=*/false,
                /*validate_full_update=*/false),
            "Failed to set cached constants for method '%s'",
            method_name.c_str());
        ET_LOG(
            Info,
            "Shared %zu cached constants into method '%s'",
            pairs.size(),
            method_name.c_str());
      }
    }

    return Error::Ok;
  }

  // Binds the whole weights blob at once, without the per-constant cache. Used
  // for artifacts whose payload names one blob rather than individual
  // constants, and it also refuses the load when that blob is needed and
  // unavailable.
  Error load_constants_legacy(
      cuda::CudaDelegateHandle* handle,
      const NamedDataMap* named_data_map,
      const std::string& weights_blob_key) const {
    // A library built before external weights cannot bind one, so there is
    // nothing to do here whether or not a blob was supplied.
    if (handle->update_constants_from_blob == nullptr) {
      ET_LOG(
          Info,
          "weights_blob '%s' is not used: this library cannot bind one",
          weights_blob_key.c_str());
      return Error::Ok;
    }

    // Fetched only once it is known to be wanted: a file-backed data map reads
    // the whole segment here, which is gigabytes for a large model.
    auto buffer_res = named_data_map->get_data(weights_blob_key.c_str());
    if (buffer_res.ok()) {
      ET_LOG(Info, "Found %s in named data map", weights_blob_key.c_str());
      const void* weights_blob = buffer_res->data();
      auto update_err = handle->update_constants_from_blob(
          handle->container_handle, static_cast<const uint8_t*>(weights_blob));
      if (update_err != Error::Ok) {
        ET_LOG(Error, "update_constants_from_blob failed");
        return update_err;
      }
      ET_CUDA_CHECK_OR_RETURN_ERROR(cudaDeviceSynchronize());
      buffer_res->Free();
    } else {
      // Without the blob the container's constant pointers stay null and the
      // failure resurfaces much later as an illegal access inside a generated
      // kernel, so report it here instead. A model with no constants needs
      // nothing bound and stays valid, which the count below decides; the count
      // is an optional symbol, so without it that cannot be established.
      ET_CHECK_OR_RETURN_ERROR(
          handle->get_num_constants != nullptr,
          NotSupported,
          "weights_blob '%s' is unavailable and this library cannot report its "
          "constant count",
          weights_blob_key.c_str());
      size_t num_constants = 0;
      ET_CHECK_OK_OR_RETURN_ERROR(
          handle->get_num_constants(handle->container_handle, &num_constants),
          "Failed to enumerate CUDA AOTI constants");
      if (num_constants != 0) {
        ET_LOG(
            Error,
            "weights_blob '%s' is unavailable, but the model has %zu constant(s) "
            "to bind",
            weights_blob_key.c_str(),
            num_constants);
        // The status from the data map, so a corrupt or unreadable sidecar is
        // not reported as an absent one.
        return buffer_res.error();
      }
      ET_LOG(
          Info,
          "weights_blob '%s' is unavailable, and the model has no constants",
          weights_blob_key.c_str());
    }
    return Error::Ok;
  }
  // Guards the singleton FQN → AtenTensorHandle cache below.
  //
  // The mutex guards init().
  // The CudaBackend instance is a process-wide singleton (registered
  // once via register_backend()), and shared_constant_tensors_ is a
  // shared-across-handles map. ExecuTorch hosts CAN call init() from
  // multiple threads when:
  //   * a multi-threaded application loads two Modules concurrently, or
  //   * a single Module is loaded from a thread pool.
  // Without the mutex, two concurrent init()s could race on
  // shared_constant_tensors_ (rehash during insert, double-insert with
  // different handles, etc.). The cost is a one-time lock during init,
  // which is negligible.
  mutable std::mutex shared_constants_mutex_;

  // FQN → AtenTensorHandle from the source (first) container.
  // The tensor handles are owned by the source container (which is never
  // explicitly deleted — see destroy() comment).
  mutable std::unordered_map<std::string, AtenTensorHandle>
      shared_constant_tensors_;

  mutable CudaWeightCache fqn_weight_cache_;
};

} // namespace executorch::backends::cuda

namespace executorch::backends {
namespace {
auto cls = cuda::CudaBackend();
executorch::runtime::Backend backend{cuda::kCudaBackendId, &cls};
static executorch::runtime::Error success_with_compiler =
    register_backend(backend);

// Auto-register the CudaAllocator so that DeviceMemoryBuffer::create(CUDA)
// works whenever the CUDA backend library is linked.
static bool cuda_allocator_registered = [] {
  executorch::runtime::register_device_allocator(
      &cuda::CudaAllocator::instance());
  return true;
}();
} // namespace
} // namespace executorch::backends