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

include/onnxruntime/ep/adapter/allocator.h

@@ -25,35 +25,47 @@ 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, which are
+  // provided and owned by the ORT framework/caller for the duration of the PrePack call (this
+  // wrapper must not release them).
+  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 +74,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 +95,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

@@ -5,6 +5,7 @@
 import os
 import tempfile
 import unittest
+from contextlib import contextmanager
 
 import numpy as np
 import torch
@@ -105,6 +106,42 @@ def is_cuda_mempool_unsupported_error(exc: Exception) -> bool:
     )
 
 
+# The fpA_intB MatMulNBits pre-pack path (which triggers the EP-provided allocator use this test
+# guards) requires a compute capability >= 7.5 (Turing) device; see MatMulNBits::MatMulNBits.
+FPA_INTB_MIN_COMPUTE_CAPABILITY = (7, 5)
+
+
+def get_cuda_compute_capability(device):
+    """Return the device (major, minor) CUDA compute capability from plugin EP metadata, or None."""
+    value = device.ep_metadata.get("cuda_compute_capability")
+    if not value:
+        return None
+    try:
+        major_str, _, minor_str = value.partition(".")
+        return (int(major_str), int(minor_str))
+    except (TypeError, ValueError):
+        return None
+
+
+def is_fpa_intb_unsupported_error(exc: Exception) -> bool:
+    """True if the error indicates the fpA_intB MatMulNBits path is unsupported for this device/config."""
+    message = str(exc)
+    return "fpA_intB" in message
+
+
+@contextmanager
+def scoped_env(name: str, value: str):
+    old_value = os.environ.get(name)
+    os.environ[name] = value
+    try:
+        yield
+    finally:
+        if old_value is None:
+            os.environ.pop(name, None)
+        else:
+            os.environ[name] = old_value
+
+
 def _create_session_options(session_config=None):
     sess_options = onnxrt.SessionOptions()
     if session_config:
@@ -220,6 +257,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 +802,62 @@ 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()
+
+        # The fpA_intB pre-pack path only runs on SM >= 7.5. Skip deterministically on older GPUs
+        # so the test does not silently pass without exercising the allocator-forwarding path.
+        compute_capability = get_cuda_compute_capability(target_device)
+        if compute_capability is None:
+            self.skipTest("CUDA plugin EP device metadata did not include cuda_compute_capability")
+        if compute_capability < FPA_INTB_MIN_COMPUTE_CAPABILITY:
+            self.skipTest(
+                "MatMulNBits fpA_intB pre-pack requires compute capability >= 7.5, but device reports "
+                f"{compute_capability[0]}.{compute_capability[1]}"
+            )
+
+        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())
+
+            # Force the fpA_intB path so weight pre-packing (MatMulNBits::PrePack_B) runs during
+            # session creation, which is where the null-allocator crash used to happen. Given the
+            # deterministic capability check above, session creation is expected to succeed, so any
+            # exception here (including the "allocator != nullptr" assertion) is a genuine regression
+            # and is deliberately propagated instead of skipped.
+            with scoped_env("ORT_FPA_INTB_GEMM", "1"):
+                sess = onnxrt.InferenceSession(model_path, sess_options=sess_options)
+
+                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)
+                # Some devices/configs may still lack a valid fpA_intB tactic at runtime; skip only
+                # for that known-unsupported case and re-raise anything else so real regressions
+                # (e.g. incorrect kernels or the allocator assertion) are not hidden.
+                try:
+                    res = sess.run(None, {"A": a})
+                except Exception as exc:
+                    if is_fpa_intb_unsupported_error(exc):
+                        self.skipTest(f"fpA_intB MatMulNBits not supported on this device: {exc}")
+                    raise
+                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):