Introduce GPUBufferPool for symmetric memory support

include/mscclpp/gpu_data_types.hpp

@@ -68,11 +68,11 @@ using __bfloat162 = __nv_bfloat162;
 
 #endif
 
-// Defined when the current device compilation target supports FP8 (`e4m3`/`e5m2`) multimem
-// instructions.
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && \
     (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__))
-#define MSCCLPP_DEVICE_FP8_MULTIMEM_SUPPORTED
+#define MSCCLPP_DEVICE_FP8_MULTIMEM_SUPPORTED 1
+#else
+#define MSCCLPP_DEVICE_FP8_MULTIMEM_SUPPORTED 0
 #endif
 
 /// Software float8 with 4 exponent bits, 3 mantissa bits, exponent bias = 15.

include/mscclpp/gpu_utils.hpp

@@ -1,10 +1,11 @@
 // Copyright (c) Microsoft Corporation.
-// Licensed under the MIT license.
+// Licensed under the MIT License.
 
 #ifndef MSCCLPP_GPU_UTILS_HPP_
 #define MSCCLPP_GPU_UTILS_HPP_
 
 #include <memory>
+#include <optional>
 #include <unordered_map>
 
 #include "env.hpp"
@@ -403,6 +404,96 @@ class GpuBuffer {
   std::shared_ptr<T> memory_;
 };
 
+namespace detail {
+
+class GpuBufferPoolStorage;
+
+}  // namespace detail
+
+/// A deterministic sub-allocation pool for GPU communication buffers.
+///
+/// The pool allocates one `GpuBuffer<char>` slab and returns reference-counted buffers for sub-ranges in that slab.
+/// If all ranks create pools with the same size and make the same allocation calls, returned buffers have identical
+/// offsets from the slab base pointer. Python bindings use this property to build symmetric memory buffers while
+/// relying on normal Python reference counting to return buffers to the local pool.
+///
+class GpuBufferPool {
+ public:
+  /// Represents one buffer returned by `GpuBufferPool`.
+  ///
+  /// This object owns a sub-range of a pool allocation. Destroying the last reference to this object returns the
+  /// sub-range to the pool. Destruction only updates local pool bookkeeping and does not perform any
+  /// cross-rank synchronization.
+  class Buffer {
+   public:
+    Buffer(const Buffer&) = delete;
+    Buffer& operator=(const Buffer&) = delete;
+    Buffer(Buffer&&) = delete;
+    Buffer& operator=(Buffer&&) = delete;
+
+    /// Destructor. Returns this buffer's sub-range to the owning pool when the last reference is destroyed.
+    ~Buffer();
+
+    /// Returns the number of bytes requested for this buffer.
+    /// @return Number of bytes in this buffer.
+    size_t bytes() const;
+
+    /// Returns this buffer's byte offset from the pool base pointer.
+    /// @return Byte offset from the pool base pointer.
+    size_t offset() const;
+
+    /// Returns the device pointer to this buffer.
+    /// @return Device pointer to this buffer.
+    char* data() const;
+
+    /// Returns the device id of the underlying pool allocation.
+    /// @return Device id of the underlying pool allocation.
+    int deviceId() const;
+
+   private:
+    friend class detail::GpuBufferPoolStorage;
+    Buffer(std::shared_ptr<detail::GpuBufferPoolStorage> storage, size_t offset, size_t bytes);
+
+    std::shared_ptr<detail::GpuBufferPoolStorage> storage_;
+    size_t offset_;
+    size_t bytes_;
+  };
+
+  /// Constructs a pool backed by a single `GpuBuffer<char>`.
+  /// @param bytes Number of bytes to reserve in the pool.
+  /// @param granularity Granularity used to size the underlying `GpuBuffer`.
+  explicit GpuBufferPool(size_t bytes, GpuBufferGranularity granularity = GpuBufferGranularity::MultiCastMinimum);
+
+  /// Allocates a sub-range from the pool.
+  /// @param bytes Number of bytes to allocate.
+  /// @param alignment Alignment in bytes for the returned offset.
+  /// @return A reference-counted pooled buffer.
+  std::shared_ptr<Buffer> allocate(size_t bytes, size_t alignment = 256);
+
+  /// Returns the number of bytes in the underlying pool allocation.
+  /// @return Number of bytes in the pool.
+  size_t bytes() const;
+
+  /// Returns the number of bytes that are available for new buffers.
+  /// @return Number of free bytes.
+  size_t freeBytes() const;
+
+  /// Returns the number of bytes currently held by active pooled buffers.
+  /// @return Number of active bytes.
+  size_t activeBytes() const;
+
+  /// Returns the device pointer to the pool base.
+  /// @return Device pointer to the pool base.
+  char* data();
+
+  /// Returns the device id of the underlying pool allocation.
+  /// @return Device id of the underlying pool allocation.
+  int deviceId() const;
+
+ private:
+  std::shared_ptr<detail::GpuBufferPoolStorage> storage_;
+};
+
 }  // namespace mscclpp
 
 #endif  // MSCCLPP_GPU_UTILS_HPP_

include/mscclpp/switch_channel_device.hpp

@@ -84,7 +84,7 @@ struct SwitchChannelDeviceHandle {
           : "l"(ptr)
           : "memory");
     }
-#if defined(MSCCLPP_DEVICE_FP8_MULTIMEM_SUPPORTED)
+#if MSCCLPP_DEVICE_FP8_MULTIMEM_SUPPORTED
     else if constexpr (std::is_same_v<VectorType, f8_e4m3x4>) {
       if constexpr (std::is_same_v<AccumT, __half>) {
         asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.e4m3x4 %0, [%1];"
@@ -198,7 +198,7 @@ struct SwitchChannelDeviceHandle {
                    "r"(val.words[1]), "r"(val.words[2]), "r"(val.words[3])
                    : "memory");
     }
-#if defined(MSCCLPP_DEVICE_FP8_MULTIMEM_SUPPORTED)
+#if MSCCLPP_DEVICE_FP8_MULTIMEM_SUPPORTED
     else if constexpr (std::is_same_v<VectorType, f8_e4m3x4>) {
       asm volatile("multimem.st.relaxed.sys.global.e4m3x4 [%0], %1;" ::"l"(ptr), "r"(val.words[0]) : "memory");
     } else if constexpr (std::is_same_v<VectorType, f8_e4m3x8>) {

python/csrc/gpu_utils_py.cpp

@@ -1,12 +1,14 @@
 // Copyright (c) Microsoft Corporation.
-// Licensed under the MIT license.
+// Licensed under the MIT License.
 
 #include <dlpack/dlpack.h>
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/shared_ptr.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/vector.h>
 
+#include <memory>
 #include <mscclpp/gpu_data_types.hpp>
 #include <mscclpp/gpu_utils.hpp>
 
@@ -15,6 +17,13 @@ using namespace mscclpp;
 
 constexpr int BYTE_BITS = 8;
 
+struct DlpackContext {
+  DLManagedTensor managedTensor{};
+  std::vector<int64_t> shape;
+  std::vector<int64_t> strides;
+  std::shared_ptr<void> owner;
+};
+
 static DLDeviceType getDeviceType() {
 #if defined(MSCCLPP_USE_ROCM)
   return kDLROCM;
@@ -52,65 +61,77 @@ static DLDataType getDlType(std::string type) {
   }
 }
 
-static nb::capsule toDlpack(GpuBuffer<char> buffer, std::string dataType, std::vector<int64_t>& shape,
-                            std::vector<int64_t>& strides) {
+static void dlpackDeleter(DLManagedTensor* self) { delete static_cast<DlpackContext*>(self->manager_ctx); }
+
+static void dlpackCapsuleDestructor(PyObject* capsule) {
+  if (PyCapsule_IsValid(capsule, "used_dltensor")) {
+    return;
+  }
+  if (!PyCapsule_IsValid(capsule, "dltensor")) {
+    return;
+  }
+  DLManagedTensor* managedTensor = static_cast<DLManagedTensor*>(PyCapsule_GetPointer(capsule, "dltensor"));
+  if (managedTensor == nullptr) {
+    return;
+  }
+  if (managedTensor->deleter) {
+    managedTensor->deleter(managedTensor);
+  }
+}
+
+static nb::capsule makeDlpack(void* data, size_t bytes, int deviceId, std::shared_ptr<void> owner, std::string dataType,
+                              std::vector<int64_t> shape, std::vector<int64_t> strides) {
   DLDataType dtype = getDlType(dataType);
-  int64_t* tensorShape = shape.size() > 0 ? new int64_t[shape.size()] : new int64_t[1];
-  int64_t* tensorStrides = strides.size() > 0 ? new int64_t[strides.size()] : nullptr;
-  if (shape.size() == 0) {
-    tensorShape[0] = (int64_t)(buffer.nelems() / ((dtype.bits * dtype.lanes + 7) / BYTE_BITS));
-  } else {
-    for (size_t i = 0; i < shape.size(); ++i) {
-      tensorShape[i] = shape[i];
+  auto ctx = std::make_unique<DlpackContext>();
+  size_t elementBytes = (dtype.bits * dtype.lanes + 7) / BYTE_BITS;
+  if (shape.empty()) {
+    if (bytes % elementBytes != 0) {
+      throw Error("DLPack buffer size must be divisible by the element size.", ErrorCode::InvalidUsage);
     }
+    ctx->shape.push_back((int64_t)(bytes / elementBytes));
+  } else {
+    ctx->shape = std::move(shape);
   }
-  for (size_t i = 0; i < strides.size(); ++i) {
-    tensorStrides[i] = strides[i];
+  ctx->strides = std::move(strides);
+  if (!ctx->strides.empty() && ctx->strides.size() != ctx->shape.size()) {
+    throw Error("DLPack strides must have the same length as shape.", ErrorCode::InvalidUsage);
   }
+  ctx->owner = std::move(owner);
 
-  DLManagedTensor* dlManagedTensor = new DLManagedTensor();
-  dlManagedTensor->dl_tensor.data = buffer.data();
+  DLManagedTensor* dlManagedTensor = &ctx->managedTensor;
+  dlManagedTensor->dl_tensor.data = data;
   dlManagedTensor->dl_tensor.device.device_type = getDeviceType();
-  dlManagedTensor->dl_tensor.device.device_id = buffer.deviceId();
-  dlManagedTensor->dl_tensor.ndim = shape.size() == 0 ? 1 : shape.size();
-  dlManagedTensor->dl_tensor.strides = tensorStrides;
-  dlManagedTensor->dl_tensor.shape = tensorShape;
+  dlManagedTensor->dl_tensor.device.device_id = deviceId;
+  dlManagedTensor->dl_tensor.ndim = static_cast<int>(ctx->shape.size());
+  dlManagedTensor->dl_tensor.strides = ctx->strides.empty() ? nullptr : ctx->strides.data();
+  dlManagedTensor->dl_tensor.shape = ctx->shape.data();
   dlManagedTensor->dl_tensor.byte_offset = 0;
   dlManagedTensor->dl_tensor.dtype = dtype;
-  dlManagedTensor->manager_ctx = new GpuBuffer<char>(buffer);
-  dlManagedTensor->deleter = [](DLManagedTensor* self) {
-    delete static_cast<GpuBuffer<char>*>(self->manager_ctx);
-    self->manager_ctx = nullptr;
-    self->dl_tensor.data = nullptr;
-    if (self->dl_tensor.shape != nullptr) {
-      delete[] self->dl_tensor.shape;
-      self->dl_tensor.shape = nullptr;
-      if (self->dl_tensor.strides) {
-        delete[] self->dl_tensor.strides;
-        self->dl_tensor.strides = nullptr;
-      }
-    }
-    delete self;
-  };
+  dlManagedTensor->manager_ctx = ctx.get();
+  dlManagedTensor->deleter = dlpackDeleter;
 
-  PyObject* dlCapsule = PyCapsule_New(static_cast<void*>(dlManagedTensor), "dltensor", [](PyObject* capsule) {
-    if (PyCapsule_IsValid(capsule, "used_dltensor")) {
-      return;
-    }
-    if (!PyCapsule_IsValid(capsule, "dltensor")) {
-      return;
-    }
-    DLManagedTensor* managedTensor = static_cast<DLManagedTensor*>(PyCapsule_GetPointer(capsule, "dltensor"));
-    if (managedTensor == nullptr) {
-      return;
-    }
-    if (managedTensor->deleter) {
-      managedTensor->deleter(managedTensor);
-    }
-  });
+  PyObject* dlCapsule = PyCapsule_New(static_cast<void*>(dlManagedTensor), "dltensor", dlpackCapsuleDestructor);
+  if (dlCapsule == nullptr) {
+    throw Error("Failed to create DLPack capsule.", ErrorCode::InvalidUsage);
+  }
+  ctx.release();
   return nb::steal<nb::capsule>(dlCapsule);
 }
 
+static nb::capsule toDlpack(GpuBuffer<char> buffer, std::string dataType, std::vector<int64_t>& shape,
+                            std::vector<int64_t>& strides) {
+  auto owner = std::make_shared<GpuBuffer<char>>(buffer);
+  return makeDlpack(buffer.data(), buffer.nelems(), buffer.deviceId(), std::move(owner), dataType, shape, strides);
+}
+
+static nb::capsule toDlpack(std::shared_ptr<GpuBufferPool::Buffer> buffer, std::string dataType,
+                            std::vector<int64_t>& shape, std::vector<int64_t>& strides) {
+  void* data = buffer->data();
+  size_t bytes = buffer->bytes();
+  int deviceId = buffer->deviceId();
+  return makeDlpack(data, bytes, deviceId, std::move(buffer), dataType, shape, strides);
+}
+
 void register_gpu_utils(nb::module_& m) {
   m.def("is_nvls_supported", &isNvlsSupported);
 
@@ -131,4 +152,25 @@ void register_gpu_utils(nb::module_& m) {
             return toDlpack(self, dataType, shape, strides);
           },
           nb::arg("data_type"), nb::arg("shape") = std::vector<int64_t>(), nb::arg("strides") = std::vector<int64_t>());
+
+  nb::class_<GpuBufferPool::Buffer>(m, "CppRawGpuBufferPoolBuffer")
+      .def("bytes", &GpuBufferPool::Buffer::bytes)
+      .def("offset", &GpuBufferPool::Buffer::offset)
+      .def("data", [](GpuBufferPool::Buffer& self) { return reinterpret_cast<uintptr_t>(self.data()); })
+      .def("device_id", &GpuBufferPool::Buffer::deviceId)
+      .def(
+          "to_dlpack",
+          [](std::shared_ptr<GpuBufferPool::Buffer> self, std::string dataType, std::vector<int64_t> shape,
+             std::vector<int64_t> strides) { return toDlpack(std::move(self), dataType, shape, strides); },
+          nb::arg("data_type"), nb::arg("shape") = std::vector<int64_t>(), nb::arg("strides") = std::vector<int64_t>());
+
+  nb::class_<GpuBufferPool>(m, "CppRawGpuBufferPool")
+      .def(nb::init<size_t, GpuBufferGranularity>(), nb::arg("bytes"),
+           nb::arg("granularity") = GpuBufferGranularity::MultiCastMinimum)
+      .def("bytes", &GpuBufferPool::bytes)
+      .def("free_bytes", &GpuBufferPool::freeBytes)
+      .def("active_bytes", &GpuBufferPool::activeBytes)
+      .def("data", [](GpuBufferPool& self) { return reinterpret_cast<uintptr_t>(self.data()); })
+      .def("device_id", &GpuBufferPool::deviceId)
+      .def("allocate", &GpuBufferPool::allocate, nb::arg("bytes"), nb::arg("alignment") = 256);
 }

python/mscclpp/__init__.py

@@ -50,6 +50,8 @@
     CppExecutionPlan as ExecutionPlan,
     CppPacketType as PacketType,
     CppRawGpuBuffer as RawGpuBuffer,
+    CppRawGpuBufferPool as RawGpuBufferPool,
+    CppRawGpuBufferPoolBuffer as RawGpuBufferPoolBuffer,
     CppReduceOp as ReduceOp,
     env,
     is_nvls_supported,
@@ -84,6 +86,8 @@
     "ExecutionPlan",
     "PacketType",
     "RawGpuBuffer",
+    "RawGpuBufferPool",
+    "RawGpuBufferPoolBuffer",
     "ReduceOp",
     "env",
     "version",
@@ -100,6 +104,7 @@
     "AlgorithmCollection",
     "CommGroup",
     "GpuBuffer",
+    "GpuBufferPool",
     "GpuBufferGranularity",
 ]