[CUDA] Fix null allocator passed to plugin EP kernel PrePack

include/onnxruntime/ep/adapter/allocator.h

@@ -25,35 +25,46 @@ class IAllocatorWrappingOrtAllocator final : public IAllocator {
  public:
   explicit IAllocatorWrappingOrtAllocator(Ort::Allocator ort_allocator)
       : IAllocator(*(EnsureOrtAllocatorHasValue(ort_allocator).GetInfo())),
-        ort_allocator_(std::move(ort_allocator)) {
+        owned_ort_allocator_(std::move(ort_allocator)),
+        ort_allocator_(owned_ort_allocator_) {
+  }
+
+  // Wraps an OrtAllocator without taking ownership. The caller must keep it alive for the
+  // lifetime of this IAllocator. This is used for allocators passed to PrePackWeight, whose
+  // lifetime is owned by the plugin EP kernel adapter.
+  explicit IAllocatorWrappingOrtAllocator(OrtAllocator* ort_allocator)
+      : IAllocator(*EnsureOrtAllocatorHasValue(ort_allocator)->Info(ort_allocator)),
+        owned_ort_allocator_(nullptr),
+        ort_allocator_(ort_allocator) {
   }
 
   void* Alloc(size_t size) override {
-    return ort_allocator_.Alloc(size);
+    return ort_allocator_->Alloc(ort_allocator_, size);
   }
 
   void Free(void* p) override {
-    ort_allocator_.Free(p);
+    ort_allocator_->Free(ort_allocator_, p);
   }
 
   void* Reserve(size_t size) override {
-    return ort_allocator_.Reserve(size);
+    if (ort_allocator_->version >= 18 && ort_allocator_->Reserve != nullptr) {
+      return ort_allocator_->Reserve(ort_allocator_, size);
+    }
+    return ort_allocator_->Alloc(ort_allocator_, size);
   }
 
   bool IsStreamAware() const override {
     static constexpr uint32_t kOrtAllocatorAllocOnStreamMinVersion = 23;
-    const OrtAllocator* raw = ort_allocator_;
-    return raw->version >= kOrtAllocatorAllocOnStreamMinVersion && raw->AllocOnStream != nullptr;
+    return ort_allocator_->version >= kOrtAllocatorAllocOnStreamMinVersion && ort_allocator_->AllocOnStream != nullptr;
   }
 
   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));
+    if (ort_allocator_->version >= kOrtAllocatorAllocOnStreamMinVersion && ort_allocator_->AllocOnStream != nullptr) {
+      return ort_allocator_->AllocOnStream(ort_allocator_, size, reinterpret_cast<OrtSyncStream*>(stream));
     }
 
-    return raw->Alloc(raw, size);
+    return ort_allocator_->Alloc(ort_allocator_, size);
   }
 
   void GetStats(AllocatorStats* stats) override {
@@ -62,10 +73,11 @@ class IAllocatorWrappingOrtAllocator final : public IAllocator {
 
     // GetStats was added in OrtAllocator version 23. For older allocators the function pointer
     // may be uninitialized, so we must not call through it.
-    const OrtAllocator* raw = ort_allocator_;
-    if (raw->version < 23 || !raw->GetStats) return;
+    if (ort_allocator_->version < 23 || !ort_allocator_->GetStats) return;
 
-    Ort::KeyValuePairs kvps = ort_allocator_.GetStats();
+    OrtKeyValuePairs* stats_kvps = nullptr;
+    Ort::ThrowOnError(ort_allocator_->GetStats(ort_allocator_, &stats_kvps));
+    Ort::KeyValuePairs kvps{stats_kvps};
     std::vector<const char*> keys, values;
     kvps.GetKeyValuePairs(keys, values);
     const size_t n = keys.size() < values.size() ? keys.size() : values.size();
@@ -82,7 +94,13 @@ class IAllocatorWrappingOrtAllocator final : public IAllocator {
     return ort_allocator;
   }
 
-  Ort::Allocator ort_allocator_;
+  static OrtAllocator* EnsureOrtAllocatorHasValue(OrtAllocator* ort_allocator) {
+    ORT_ENFORCE(ort_allocator != nullptr, "OrtAllocator must be non-null.");
+    return ort_allocator;
+  }
+
+  Ort::Allocator owned_ort_allocator_;
+  OrtAllocator* ort_allocator_{};
 
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(IAllocatorWrappingOrtAllocator);
 };

include/onnxruntime/ep/adapter/op_kernel.h

@@ -9,6 +9,7 @@
 
 #include <memory>
 #include <type_traits>
+#include <utility>
 #include <vector>
 
 #include "core/framework/allocator.h"
@@ -222,14 +223,15 @@ struct KernelImpl : OrtKernelImpl {
   static OrtStatus* ORT_API_CALL PrePackWeightImpl(_In_ OrtKernelImpl* this_ptr,
                                                    _In_ const OrtValue* weight,
                                                    int input_index,
-                                                   _In_ OrtAllocator* /* allocator */,
+                                                   _In_ OrtAllocator* allocator,
                                                    _In_opt_ OrtSharedPrePackedWeightCache* /* prepacked_weight_cache */,
                                                    _Out_ bool* is_packed) noexcept {
     Status status;
     ORT_TRY {
       auto* kernel_impl = static_cast<KernelImpl*>(this_ptr)->impl_.get();
       const auto tensor = CreateTensorFromApiValue(const_cast<OrtValue*>(weight));
-      status = kernel_impl->PrePack(tensor, input_index, AllocatorPtr{}, *is_packed, nullptr);
+      auto allocator_wrapper = std::make_shared<IAllocatorWrappingOrtAllocator>(allocator);
+      status = kernel_impl->PrePack(tensor, input_index, std::move(allocator_wrapper), *is_packed, nullptr);
     }
     ORT_CATCH(const std::exception& ex) {
       ORT_HANDLE_EXCEPTION([&]() {

onnxruntime/test/python/transformers/test_cuda_plugin_ep.py

@@ -220,6 +220,52 @@ def create_gemm_model(model_path, alpha=1.0, beta=1.0, transA=0, transB=0):
     save(model_def, model_path)
 
 
+def create_matmul_nbits_model(model_path, k=64, n=64, bits=4, block_size=32):
+    # Create a MatMulNBits (com.microsoft) model with a runtime-prepacked (weight_prepacked=0)
+    # fp16 quantized weight. During session initialization the CUDA kernel pre-packs the constant
+    # weight via MatMulNBits::PrePack_B, which allocates scratch through the EP-provided allocator
+    # (IAllocator::MakeUniquePtr(alloc, ...)). This is a regression guard: the CUDA-EP-as-plugin
+    # adapter previously forwarded a null allocator into PrePack, so this model crashed at session
+    # creation with "IAllocator::ValidateAllocator ... allocator != nullptr was false".
+    k_blocks = (k + block_size - 1) // block_size
+    blob_size = block_size * bits // 8
+
+    # Deterministic quantized weight and positive scales. Exact numerics are not asserted here;
+    # the test only requires that pre-packing runs (and no longer crashes) end to end.
+    b = np.full((n, k_blocks, blob_size), 0x88, dtype=np.uint8)
+    scales = np.full((n, k_blocks), 0.02, dtype=np.float16)
+
+    node_def = helper.make_node(
+        "MatMulNBits",
+        ["A", "B", "scales"],
+        ["Y"],
+        domain="com.microsoft",
+        K=k,
+        N=n,
+        bits=bits,
+        block_size=block_size,
+        weight_prepacked=0,
+    )
+    graph_def = helper.make_graph(
+        [node_def],
+        "test-model-matmulnbits",
+        [helper.make_tensor_value_info("A", TensorProto.FLOAT16, [2, k])],
+        [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [2, n])],
+        initializer=[
+            helper.make_tensor("B", TensorProto.UINT8, b.shape, b.tobytes(), raw=True),
+            helper.make_tensor("scales", TensorProto.FLOAT16, scales.shape, scales.tobytes(), raw=True),
+        ],
+    )
+    opset_onnx = OperatorSetIdProto()
+    opset_onnx.version = 21
+    opset_ms = OperatorSetIdProto()
+    opset_ms.domain = "com.microsoft"
+    opset_ms.version = 1
+    model_def = helper.make_model(graph_def, producer_name="onnx-example", opset_imports=[opset_onnx, opset_ms])
+    model_def.ir_version = 10
+    save(model_def, model_path)
+
+
 def create_conv_model(model_path):
     # Create a simple Conv model: Y = Conv(X, W)
     node_def = helper.make_node("Conv", ["X", "W"], ["Y"], pads=[1, 1, 1, 1], strides=[1, 1], dilations=[1, 1], group=1)
@@ -719,6 +765,45 @@ def test_registration_conv(self):
         result = run_operator_test(target_device, create_conv_model, inputs, _expected_conv)
         self.assertTrue(result, "Conv plugin registration test failed")
 
+    def test_registration_matmul_nbits_prepack(self):
+        # Regression guard for the CUDA-EP-as-plugin pre-pack allocator bug: a fp16 MatMulNBits
+        # weight is pre-packed during session initialization, which allocates scratch through the
+        # EP-provided allocator. The plugin op-kernel adapter previously forwarded a null allocator
+        # into PrePack, crashing session creation with an IAllocator::ValidateAllocator failure.
+        target_device = get_cuda_plugin_device()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmp:
+            model_path = tmp.name
+        try:
+            create_matmul_nbits_model(model_path)
+            sess_options = _create_session_options()
+            sess_options.add_provider_for_devices([target_device], _plugin_provider_options())
+
+            # Session creation is where weight pre-packing runs (and where the null-allocator crash
+            # used to happen). If the device lacks fpA_intB GEMM support the op is not pre-packed and
+            # falls back, so treat an unsupported-op error as a skip rather than a failure.
+            try:
+                sess = onnxrt.InferenceSession(model_path, sess_options=sess_options)
+            except Exception as exc:
+                if "allocator != nullptr" in str(exc):
+                    raise
+                self.skipTest(f"MatMulNBits pre-pack not supported on this device: {exc}")
+            else:
+                assigned_nodes, assignment_info = _get_assigned_nodes(sess, CUDA_PLUGIN_EP_NAME)
+                self.assertTrue(
+                    assigned_nodes,
+                    f"{CUDA_PLUGIN_EP_NAME} was assigned no nodes. "
+                    f"Assignments: {_format_assignment_summary(assignment_info)}",
+                )
+
+                a = np.random.rand(2, 64).astype(np.float16)
+                res = sess.run(None, {"A": a})
+                self.assertEqual(res[0].shape, (2, 64))
+                self.assertTrue(np.isfinite(res[0].astype(np.float32)).all(), "MatMulNBits produced non-finite output")
+        finally:
+            if os.path.exists(model_path):
+                os.remove(model_path)
+
     # ---- Provider options tests ----
 
     def test_provider_options_valid(self):