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,11 +1,12 @@
 // 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 <unordered_map>
+#include <utility>
 
 #include "env.hpp"
 #include "errors.hpp"
@@ -143,6 +144,8 @@ std::shared_ptr<GpuStreamPool> gpuStreamPool();
 
 namespace detail {
 
+class GpuBufferPoolStorage;
+
 void setReadWriteMemoryAccess(void* base, size_t size);
 int gpuIdFromAddress(void* ptr);
 
@@ -332,7 +335,8 @@ enum class GpuBufferGranularity {
   MultiCastRecommended,
 };
 
-/// Allocates a GPU memory space specialized for communication. The memory is zeroed out. Get the device pointer by
+/// Owns a GPU memory space specialized for communication. Memory allocated by the public constructor is zeroed.
+/// Memory returned by `GpuBufferPool::allocate()` may contain data from a previous use. Get the device pointer by
 /// `GpuBuffer::data()`.
 ///
 /// Use this function for communication buffers, i.e., only when other devices (CPU, GPU, NIC, etc.) may access this
@@ -402,12 +406,63 @@ class GpuBuffer {
   int deviceId() const { return deviceId_; }
 
  private:
+  friend class detail::GpuBufferPoolStorage;
+
+  GpuBuffer(size_t nelems, size_t bytes, int deviceId, std::shared_ptr<T> memory)
+      : nelems_(nelems), bytes_(bytes), deviceId_(deviceId), memory_(std::move(memory)) {}
+
   size_t nelems_;
   size_t bytes_;
   int deviceId_;
   std::shared_ptr<T> memory_;
 };
 
+/// 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.
+/// Callers must keep each returned buffer alive until all GPU work using it has completed.
+///
+class GpuBufferPool {
+ public:
+  /// 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.
+  /// The returned memory has unspecified contents and may contain data from a previous allocation.
+  /// @param bytes Number of bytes to allocate.
+  /// @param alignment Alignment in bytes for the returned offset.
+  /// @return A reference-counted `GpuBuffer<char>` backed by the pool.
+  std::shared_ptr<GpuBuffer<char>> 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

@@ -110,7 +110,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];"
@@ -224,7 +224,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,67 @@ 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 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 = [](DLManagedTensor* self) { delete static_cast<DlpackContext*>(self->manager_ctx); };
 
-  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);
+}
+
 void register_gpu_utils(nb::module_& m) {
   m.def("is_nvls_supported", &isNvlsSupported);
   m.def("is_bulk_supported", &isBulkSupported);
@@ -132,4 +143,14 @@ 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>(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

@@ -102,6 +102,7 @@
     "AlgorithmCollection",
     "CommGroup",
     "GpuBuffer",
+    "GpuBufferPool",
     "GpuBufferGranularity",
 ]
 

python/mscclpp/_core/buffer.py

@@ -6,9 +6,13 @@
 
 import cupy as cp
 import numpy as np
-from mscclpp._mscclpp import CppRawGpuBuffer, CppGpuBufferGranularity
+from mscclpp._mscclpp import (
+    CppRawGpuBuffer,
+    CppRawGpuBufferPool,
+    CppGpuBufferGranularity,
+)
 
-__all__ = ["GpuBuffer", "GpuBufferGranularity"]
+__all__ = ["GpuBuffer", "GpuBufferPool", "GpuBufferGranularity"]
 
 GpuBufferGranularity = CppGpuBufferGranularity
 
@@ -35,3 +39,65 @@ def __new__(
         buffer = CppRawGpuBuffer(np.prod(shape) * np.dtype(dtype).itemsize, granularity)
         memptr = cp.cuda.MemoryPointer(cp.cuda.UnownedMemory(buffer.data(), buffer.bytes(), buffer), 0)
         return cp.ndarray(shape, dtype=dtype, strides=strides, order=order, memptr=memptr)
+
+
+class GpuBufferPool:
+    """A GPU buffer pool that returns raw buffers backed by one communication-friendly allocation.
+
+    All ranks should create the same-sized pool and call :meth:`allocate` in the same order to get matching offsets.
+    Keep each returned buffer alive until all GPU work using it has completed.
+    """
+
+    def __init__(
+        self,
+        nbytes: int,
+        granularity: CppGpuBufferGranularity = CppGpuBufferGranularity.MultiCastMinimum,
+    ):
+        if nbytes <= 0:
+            raise ValueError("Pool size must be positive.")
+        self._pool = CppRawGpuBufferPool(int(nbytes), granularity)
+
+    @property
+    def bytes(self) -> int:
+        """Number of bytes in the underlying pool allocation."""
+        return self._pool.bytes()
+
+    @property
+    def free_bytes(self) -> int:
+        """Number of bytes available for new buffers."""
+        return self._pool.free_bytes()
+
+    @property
+    def active_bytes(self) -> int:
+        """Number of bytes currently held by live raw buffers."""
+        return self._pool.active_bytes()
+
+    @property
+    def data(self) -> int:
+        """Device pointer to the pool base allocation."""
+        return self._pool.data()
+
+    @property
+    def device_id(self) -> int:
+        """CUDA/HIP device id of the pool allocation."""
+        return self._pool.device_id()
+
+    def allocate(
+        self,
+        nbytes: int,
+        alignment: int = 256,
+    ) -> CppRawGpuBuffer:
+        """Allocate a raw buffer from the pool.
+
+        Args:
+            nbytes: Number of bytes to allocate.
+            alignment: Required byte alignment of the allocation offset from the pool base.
+
+        The returned memory has unspecified contents and may contain data from a previous allocation.
+        """
+        if nbytes <= 0:
+            raise ValueError("Buffer size must be positive.")
+        if alignment <= 0:
+            raise ValueError("Alignment must be positive.")
+
+        return self._pool.allocate(int(nbytes), int(alignment))