[CUDA] Support user compute stream with CUDA graph in CUDA plugin EP

cmake/onnxruntime_providers_cuda_plugin.cmake

@@ -88,10 +88,11 @@ list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX ".*/tensor/sequence_op\\.cc$")
 # in the CPU provider and is not linked into the plugin.
 list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX ".*/tensor/size\\.cc$")
 
-# Permanently excluded — pure CPU ops, handled by GetCpuPreferredNodes.
-# shape_op.cc inherits from onnxruntime::OpKernel (framework)
-# which cannot convert to ep::adapter::OpKernel in the plugin build.
-list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX ".*/tensor/shape_op\\.cc$")
+# shape_op.cc is INCLUDED in the plugin build. It provides an adapter-based
+# Shape kernel under #ifdef BUILD_CUDA_EP_AS_PLUGIN (the CPU onnxruntime::Shape
+# class, which derives from the framework OpKernel, is only used in the
+# non-plugin build). Registering Shape on the EP keeps it off the CPU EP and
+# avoids Memcpy nodes that would otherwise break CUDA Graph capture.
 
 # Exclude contrib training ops (shrunken_gather depends on provider_api.h in header).
 list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX ".*/contrib_ops/cuda/tensor/shrunken_gather\\.cc$")

include/onnxruntime/ep/adapter/allocator.h

@@ -41,21 +41,19 @@ class IAllocatorWrappingOrtAllocator final : public IAllocator {
   }
 
   bool IsStreamAware() const override {
-    return false;
-
-    // TODO: Enable once AllocOnStream() is implemented.
-    // static constexpr uint32_t kOrtAllocatorAllocOnStreamMinVersion = 23;
-    // const OrtAllocator* raw = ort_allocator_;
-    // return raw->version >= kOrtAllocatorAllocOnStreamMinVersion && raw->AllocOnStream != nullptr;
+    static constexpr uint32_t kOrtAllocatorAllocOnStreamMinVersion = 23;
+    const OrtAllocator* raw = ort_allocator_;
+    return raw->version >= kOrtAllocatorAllocOnStreamMinVersion && raw->AllocOnStream != nullptr;
   }
 
-  void* AllocOnStream(size_t /*size*/, Stream* /*stream*/) override {
-    // TODO: Implement AllocOnStream().
-    // The internal `onnxruntime::IAllocator::AllocOnStream` signature takes an internal `onnxruntime::Stream*`
-    // argument, while the public `::OrtAllocator::AllocOnStream` signature takes an `::OrtSyncStream*` argument.
-    // We need to properly map from one to the other.
-    // `::OrtSyncStream*` should be treated as an opaque type from the plugin EP's perspective.
-    ORT_NOT_IMPLEMENTED("IAllocatorWrappingOrtAllocator::AllocOnStream is not implemented yet.");
+  void* AllocOnStream(size_t size, Stream* stream) override {
+    static constexpr uint32_t kOrtAllocatorAllocOnStreamMinVersion = 23;
+    OrtAllocator* raw = ort_allocator_;
+    if (raw->version >= kOrtAllocatorAllocOnStreamMinVersion && raw->AllocOnStream != nullptr) {
+      return raw->AllocOnStream(raw, size, reinterpret_cast<OrtSyncStream*>(stream));
+    }
+
+    return raw->Alloc(raw, size);
   }
 
   void GetStats(AllocatorStats* stats) override {

onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc

@@ -152,6 +152,16 @@ Status GroupQueryAttention<T, U>::PrePack(const Tensor& tensor, int input_idx, A
   // XQA consumes the attention sink as FP32. When head_sink is a constant initializer, convert it once
   // here into a cached device buffer (xqa_head_sink_) to avoid a per-launch conversion. Dynamic /
   // non-initializer head_sink inputs are not prepacked and fall back to the per-launch scratch path.
+  // In the plugin EP build the AllocatorPtr passed by the framework can arrive null across the
+  // library boundary. Fall back to the kernel's own default-memory allocator, which is always valid.
+  // This kernel keeps is_packed=false and never populates prepacked_weights, so the framework does
+  // not register a PrePackedWeights container for this input and there is no disk externalization of
+  // the result. The allocator is used only to materialize the cached FP32 sink in device memory (an
+  // H2D copy followed by a conversion kernel), so a default-memory (device) allocator is exactly the
+  // kind required here -- a CPU allocator would be wrong.
+  if (!alloc) {
+    alloc = this->Info().GetAllocator(OrtMemType::OrtMemTypeDefault);
+  }
   if constexpr (std::is_same_v<T, MLFloat16> || std::is_same_v<T, BFloat16>) {
     const auto& shape = tensor.Shape();
     ORT_RETURN_IF_NOT(shape.NumDimensions() == 1,

onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc

@@ -1064,6 +1064,12 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
   ORT_UNUSED_PARAMETER(prepacked_weights);
   is_packed = false;
 
+  // In the plugin EP build the AllocatorPtr passed by the framework can arrive null across the
+  // library boundary. Fall back to the kernel's own default-memory allocator, which is always valid.
+  if (!alloc) {
+    alloc = this->Info().GetAllocator(OrtMemType::OrtMemTypeDefault);
+  }
+
   cudaStream_t stream = 0;  // Use default stream for PrePack operations
 
   DUMP_TENSOR_INIT();

onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc

@@ -105,6 +105,11 @@ Status MatMulNBits<T>::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr
                                /*out*/ bool& is_packed,
                                /*out*/ PrePackedWeights* /*prepacked_weights*/) {
   is_packed = false;
+  // In the plugin EP build the AllocatorPtr passed by the framework can arrive null across the
+  // library boundary. Fall back to the kernel's own default-memory allocator, which is always valid.
+  if (!alloc) {
+    alloc = this->Info().GetAllocator(OrtMemType::OrtMemTypeDefault);
+  }
   if constexpr (std::is_same_v<T, MLFloat16> || std::is_same_v<T, BFloat16>) {
     if (has_fpA_intB_gemm_) {
       cudaStream_t stream = cudaStreamLegacy;  // Use default stream for prepacking.

onnxruntime/core/framework/session_state.cc

@@ -603,7 +603,8 @@ Status SessionState::PrepackConstantInitializedTensors(
                   // within this session. Or if the weight is not present on disk,
                   // we store the newly minted pre-packed data.
 
-                  AllocatorPtr session_initializer_alloc = GetInitializerAllocator(kernel->Info().GetDevice(OrtMemType::OrtMemTypeDefault));
+                  AllocatorPtr session_initializer_alloc = GetInitializerAllocator(
+                      kernel->Info().GetDevice(OrtMemType::OrtMemTypeDefault));
                   PrePackedWeights weights_to_be_filled_in;
                   // The reason we invoke PrePack() before looking into the container for any pre-packed weight
                   // cached by another instance of the same op_type (for the same constant initializer) is because

onnxruntime/core/providers/cuda/plugin/cuda_ep.cc

@@ -83,20 +83,29 @@ void DestroyCudaStreamForDevice(cudaStream_t stream, int device_id) {
 }  // namespace
 
 struct CudaEp::PerThreadContext {
-  explicit PerThreadContext(int device_id)
+  // When external_graph_stream is non-null (user_compute_stream combined with CUDA graph),
+  // capture and replay happen on that user-owned stream so they see the same stream as the
+  // kernels; the context neither creates nor destroys it. Otherwise the context creates and
+  // owns a dedicated graph stream.
+  explicit PerThreadContext(int device_id, cudaStream_t external_graph_stream = nullptr)
       : device_id(device_id),
-        graph_stream(CreateCudaStreamForDevice(device_id)),
+        owns_graph_stream(external_graph_stream == nullptr),
+        graph_stream(external_graph_stream != nullptr ? external_graph_stream
+                                                      : CreateCudaStreamForDevice(device_id)),
         cuda_graph(graph_stream) {
   }
 
   ~PerThreadContext() {
     // Destroy captured graph execs before destroying the stream they replay on.
     cuda_graph.Reset();
-    DestroyCudaStreamForDevice(graph_stream, device_id);
+    if (owns_graph_stream) {
+      DestroyCudaStreamForDevice(graph_stream, device_id);
+    }
     graph_stream = nullptr;
   }
 
   int device_id;
+  bool owns_graph_stream;
   cudaStream_t graph_stream = nullptr;
   CudaGraphManager cuda_graph;
   size_t pre_capture_free_mem = 0;
@@ -392,7 +401,9 @@ OrtStatus* ORT_API_CALL CudaEp::CreateSyncStreamForDeviceImpl(
   auto cuda_stream = std::make_unique<CudaSyncStream>(ep->factory_, device_id, this_ptr);
 
   if (ep->config_.has_user_compute_stream && ep->config_.user_compute_stream != nullptr) {
-    // Wrap the user-provided external CUDA stream with full cuBLAS/cuDNN handles.
+    // Wrap the user-provided external CUDA stream with full cuBLAS/cuDNN handles. When CUDA
+    // graph capture is also enabled, capture/replay run on this same user stream (see
+    // GetPerThreadContext), so kernels and graph capture share one stream.
     RETURN_IF_ERROR(cuda_stream->InitHandlesWithUserStream(
         static_cast<cudaStream_t>(ep->config_.user_compute_stream)));
   } else if (ep->config_.enable_cuda_graph) {
@@ -467,7 +478,14 @@ CudaEp::PerThreadContext& CudaEp::GetPerThreadContext() const {
     return *cached_context_it->second;
   }
 
-  auto context = std::make_shared<PerThreadContext>(config_.device_id);
+  // When a user compute stream is combined with CUDA graph capture, capture/replay must run on
+  // the user's stream (the same stream the kernels are issued to) rather than a separate
+  // EP-owned stream. The user owns the stream's lifetime, so the context must not destroy it.
+  cudaStream_t external_graph_stream =
+      (config_.has_user_compute_stream && config_.enable_cuda_graph)
+          ? static_cast<cudaStream_t>(config_.user_compute_stream)
+          : nullptr;
+  auto context = std::make_shared<PerThreadContext>(config_.device_id, external_graph_stream);
   PerThreadContext& context_ref = *context;
   {
     std::lock_guard<std::mutex> lock(per_thread_contexts_mutex_);

onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc

@@ -609,12 +609,9 @@ OrtStatus* ORT_API_CALL CudaEpFactory::CreateEpImpl(
         "CUDA plugin EP does not support using both user_compute_stream and external allocator simultaneously.");
   }
 
-  // Validate: user_compute_stream and cuda graph cannot both be active.
-  if (config.has_user_compute_stream && config.enable_cuda_graph) {
-    return factory->ort_api_.CreateStatus(
-        ORT_INVALID_ARGUMENT,
-        "CUDA plugin EP does not support using both user_compute_stream and enable_cuda_graph simultaneously.");
-  }
+  // user_compute_stream and enable_cuda_graph CAN be combined: when both are set, CUDA graph
+  // capture/replay runs on the user-provided stream (the same stream kernels are issued to),
+  // matching the bundled CUDA EP behavior. See CudaEp::GetPerThreadContext.
 
   // When user_compute_stream is set, force unified stream mode (matches bundled EP behavior).
   if (config.has_user_compute_stream) {

onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h

@@ -1025,54 +1025,53 @@ class CudaKernel : public OpKernel {
   template <typename T>
   inline IAllocatorUniquePtr<T> GetScratchBuffer(size_t cnt, void* s) const {
     if (cnt == 0) return IAllocatorUniquePtr<T>(nullptr, [](T*) {});
-    size_t sz = 0;
-    if (!detail::TryBytesForCount(cnt, detail::SizeOf<T>::value, sz)) {
-      ORT_THROW("CUDA scratch buffer allocation size overflow for ", cnt, " elements");
-    }
-
-    void* p = nullptr;
-    cudaError_t alloc_result = cudaSuccess;
-    bool used_async_alloc = false;
-    if (s) {
-      // Note: stream-ordered allocations (cudaMallocAsync/cudaFreeAsync) rely on CUDA Memory Pools,
-      // which are not supported on NVIDIA GPUs with Multi-Instance GPU (MIG) enabled.
-      // On such instances, this will return cudaErrorNotSupported.
-      alloc_result = cudaMallocAsync(&p, sz, static_cast<cudaStream_t>(s));
-      used_async_alloc = (alloc_result == cudaSuccess);
-      if (!used_async_alloc && (alloc_result == cudaErrorNotSupported || alloc_result == cudaErrorInvalidValue)) {
-        cudaGetLastError();  // Clear the thread-local error state
-        alloc_result = cudaMalloc(&p, sz);
-      }
-    } else {
-      alloc_result = cudaMalloc(&p, sz);
-    }
 
-    if (alloc_result != cudaSuccess) {
-      ORT_THROW("CUDA scratch buffer allocation failed for ", sz, " bytes: ", cudaGetErrorString(alloc_result));
+    // Route kernel scratch/workspace allocations through the EP allocator
+    // (a BFC arena by default) instead of raw cudaMallocAsync/cudaMalloc.
+    //
+    // The arena pre-reserves device memory and reuses freed chunks across runs.
+    // Once the model has executed `min_num_runs_before_cuda_graph_capture`
+    // warmup runs, the arena has grown to its steady-state working set, so the
+    // capture run serves every scratch allocation from an already-reserved chunk
+    // without issuing a fresh cudaMalloc. This keeps the device free-memory
+    // footprint stable across the capture window, which is required for correct
+    // CUDA graph capture/replay.
+    //
+    // The previous behavior (cudaMallocAsync/cudaMalloc allocated-and-freed per
+    // call) allocated new device memory on every run, including the capture run,
+    // so no amount of warmup could stabilize it and the
+    // "GPU memory was allocated during CUDA graph capture" detector would trip.
+    // This now matches the built-in (non-plugin) CUDA EP, which also obtains
+    // scratch from Info().GetAllocator() (see core/providers/cuda/cuda_kernel.h).
+    //
+    // Forward the compute stream so the allocation (and any stream-aware free)
+    // is ordered on the same stream the kernel enqueues work on. The plugin
+    // device arena is stream-aware: it tracks chunk-to-stream assignments and,
+    // for a CUDA mempool arena, issues cudaMallocFromPoolAsync/cudaFreeAsync on
+    // the allocation stream. Allocating/freeing on a different stream than the
+    // consuming kernel can race and lead to use-after-free. This also matches
+    // the built-in CUDA EP, which forwards its compute stream to MakeUniquePtr.
+    // The overflow check that the previous hand-rolled path performed is still
+    // enforced inside MakeUniquePtr via ValidatedCalcMemSizeForArray (it throws
+    // on cnt * sizeof(T) overflow).
+#ifndef __CUDACC__
+    onnxruntime::OrtStreamAdapter ort_stream(s);
+    return ::onnxruntime::IAllocator::MakeUniquePtr<T>(
+        Info().GetAllocator(OrtMemType::OrtMemTypeDefault), cnt, /*use_reserve*/ false,
+        ort_stream.get());
+#else
+    // The framework Stream type is incomplete in NVCC translation units, so the
+    // PluginStreamShim that bridges a raw cudaStream_t to a framework Stream*
+    // cannot be constructed here. Do not silently drop a non-null stream:
+    // stream-aware allocators rely on the allocation stream for ordered reuse/free.
+    if (s != nullptr) {
+      ORT_THROW("Plugin CUDA GetScratchBuffer cannot allocate with a non-null stream from an NVCC "
+                "translation unit.");
     }
-
-    return IAllocatorUniquePtr<T>(static_cast<T*>(p), [s, used_async_alloc](T* ptr) {
-      if (ptr) {
-        // Guard: only attempt async free if the stream is still registered.
-        // CudaSyncStream::~CudaSyncStream guarantees UnregisterStream() is
-        // called before cudaStreamDestroy(), so a non-null lookup here means
-        // the raw cudaStream_t handle is still valid.
-        if (used_async_alloc && s &&
-            cuda_plugin::CudaSyncStream::FromCudaStream(static_cast<cudaStream_t>(s)) != nullptr) {
-          // As noted above, cudaFreeAsync may also return cudaErrorNotSupported on MIG-enabled instances.
-          cudaError_t free_result = cudaFreeAsync(ptr, static_cast<cudaStream_t>(s));
-          if (free_result == cudaSuccess) {
-            return;
-          }
-          cudaGetLastError();  // Clear any error set by cudaFreeAsync
-        }
-
-        // Fall back to synchronous free if async free is unsupported or if the
-        // stream is no longer registered. cudaFree is valid for allocations
-        // returned by cudaMallocAsync and avoids using a stale stream handle.
-        cudaFree(ptr);
-      }
-    });
+    return ::onnxruntime::IAllocator::MakeUniquePtr<T>(
+        Info().GetAllocator(OrtMemType::OrtMemTypeDefault), cnt, /*use_reserve*/ false,
+        /*stream*/ nullptr);
+#endif
   }
   template <typename T>
   inline IAllocatorUniquePtr<T> GetTransientScratchBuffer(size_t cnt) const {

onnxruntime/core/providers/cuda/tensor/shape_op.cc

@@ -2,12 +2,87 @@
 // Licensed under the MIT License.
 
 #include "core/providers/shared_library/provider_api.h"
+#ifndef BUILD_CUDA_EP_AS_PLUGIN
 #include "core/providers/cpu/tensor/shape_op.h"
+#endif
 #include "core/providers/cuda/cuda_fwd.h"
 
+#ifdef BUILD_CUDA_EP_AS_PLUGIN
+#include <limits>
+#endif
+
 namespace onnxruntime {
 namespace cuda {
 
+#ifdef BUILD_CUDA_EP_AS_PLUGIN
+// The bundled CUDA EP registers the CPU `onnxruntime::Shape` kernel (which
+// derives from the framework `OpKernel`) and only marks its output as CPU
+// memory. That class cannot be registered through the plugin EP's adapter
+// kernel machinery, so the plugin build provides an adapter-based Shape kernel
+// with identical semantics. Shape only reads the input's shape metadata (never
+// its data) and writes the dims to a CPU output, so registering it on the CUDA
+// EP keeps the node inside the device partition and avoids the device->host
+// Memcpy node that the framework would otherwise insert to feed an isolated CPU
+// Shape node -- a Memcpy that would prevent CUDA Graph capture. The output is
+// still CPU memory, so a downstream device consumer may still need a copy; this
+// removes the graph-breaking input-side Memcpy, it does not eliminate all copies.
+class Shape final : public CudaKernel {
+ public:
+  explicit Shape(const OpKernelInfo& info) : CudaKernel(info) {
+    info.GetAttrOrDefault<int64_t>("start", &start_index_, 0);
+
+    if (start_index_ != 0) {
+      // "start" is provided and is non-default (default is 0)
+      needs_slicing_ = true;
+    }
+
+    if (info.GetAttr<int64_t>("end", &end_index_).IsOK()) {
+      needs_slicing_ = true;
+    }
+  }
+
+  // Takes a tensor as input and outputs a 1D int64 tensor (on CPU memory)
+  // containing the shape of the input tensor.
+  Status ComputeInternal(OpKernelContext* context) const override {
+    const auto* input = context->Input<Tensor>(0);
+    const TensorShape& input_shape = input->Shape();
+
+    int64_t rank = static_cast<int64_t>(input_shape.NumDimensions());
+
+    if (!needs_slicing_) {  // vanilla use of Shape (no slicing)
+      Tensor* output = context->Output(0, {rank});
+      input_shape.CopyDims(output->MutableData<int64_t>(), static_cast<size_t>(rank));
+    } else {  // slicing is needed
+      int64_t true_start = start_index_;
+      int64_t true_end = end_index_;
+
+      // Deal with negative(s) and clamp
+      true_start = true_start < 0 ? true_start + rank : true_start;
+      true_start = true_start < 0 ? 0 : ((true_start > rank) ? rank : true_start);
+
+      true_end = true_end < 0 ? true_end + rank : true_end;
+      true_end = true_end < 0 ? 0 : ((true_end > rank) ? rank : true_end);
+
+      auto slice_length = true_end - true_start;
+      Tensor* output = context->Output(0, {slice_length < 0 ? 0 : slice_length});
+
+      if (slice_length > 0) {
+        input_shape.CopyDims(output->MutableData<int64_t>(),
+                             static_cast<size_t>(true_start),
+                             static_cast<size_t>(slice_length));
+      }
+    }
+
+    return Status::OK();
+  }
+
+ private:
+  bool needs_slicing_ = false;
+  int64_t start_index_ = 0;
+  int64_t end_index_ = std::numeric_limits<int64_t>::max();
+};
+#endif  // BUILD_CUDA_EP_AS_PLUGIN
+
 ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     Shape,
     kOnnxDomain,