cpp/tensorrt_llm/runtime/mcastDeviceMemory.cpp (3)

173-180: Missing error checks for CUDA host allocations and CU memset

cudaMallocHost/cudaFreeHost and cuMemsetD8 return status; unchecked failures can corrupt flow.

Apply:

-    cudaMallocHost(&fabric_handle, sizeof(CUmemFabricHandle));
+    TLLM_CUDA_CHECK(cudaMallocHost(reinterpret_cast<void**>(&fabric_handle), sizeof(CUmemFabricHandle)));
@@
-    cudaFreeHost(fabric_handle);
+    TLLM_CUDA_CHECK(cudaFreeHost(fabric_handle));

And above:

-    cudaMallocHost(&exphndl, mGroupSize * sizeof(CUmemFabricHandle));
+    TLLM_CUDA_CHECK(cudaMallocHost(reinterpret_cast<void**>(&exphndl), mGroupSize * sizeof(CUmemFabricHandle)));
@@
-    cudaFreeHost(exphndl);
+    TLLM_CUDA_CHECK(cudaFreeHost(exphndl));

Additionally:

-            cuMemsetD8(mSignalPads[i], 0, kSIGNAL_PAD_SIZE);
+            TLLM_CU_CHECK(cuMemsetD8(mSignalPads[i], 0, kSIGNAL_PAD_SIZE));

---

142-149: roundUp assumes power-of-two granularity; use general formula

alloc_granularity and mc_granularity are not guaranteed to be powers of two. Bit-masking can under/over-round.

Apply:

-inline size_t roundUp(size_t val, size_t gran)
-{
-    return (val + gran - 1) & ~(gran - 1);
-}
+inline size_t roundUp(size_t val, size_t gran)
+{
+    return ((val + gran - 1) / gran) * gran;
+}

---

113-124: Leak: address space reserved for UC region is never freed

allocMnMcastMem reserves a contiguous VA range (ptr) for all UC maps but only unmaps each slice; it never calls cuMemAddressFree on the UC base. This leaks VA space until process exit.

Apply:

     if (mIsMNNvlink)
     {
         for (uint32_t rank = 0; rank < mGroupSize; rank++)
         {
             TLLM_CU_CHECK(cuMemUnmap(mUcPtrs[rank], mAllocationSize));
             // We need to release the handle on each rank
             TLLM_CU_CHECK(cuMemRelease(mUcHandles[rank]));
         }
+        if (!mUcPtrs.empty())
+        {
+            TLLM_CU_CHECK(cuMemAddressFree(mUcPtrs[0], mAllocationSize * mGroupSize));
+        }
         TLLM_CU_CHECK(cuMemUnmap(mMcPtr, mAllocationSize));
         TLLM_CU_CHECK(cuMemAddressFree(mMcPtr, mAllocationSize));
         TLLM_CU_CHECK(cuMemRelease(mMcHandle));
     }

cpp/tensorrt_llm/runtime/mcastGPUBuffer.h (2)

54-68: OOB risk: negative storageOffset and size overflow not handled; shape type should be int64_t.

• storageOffset < 0 will underflow when converted to size_t and can bypass the bound check.
• (numel + storageOffset) * elementSize may overflow size_t before the check.
• Prefer std::vector<int64_t> (PyTorch convention) and pass by const& to avoid copies.
• Error identifier casing should match the class/method for easier grep.

Apply:

-    at::Tensor getUCBuffer(uint32_t rank, std::vector<long int> sizes, torch::ScalarType dtype, int64_t storageOffset)
+    at::Tensor getUCBuffer(
+        uint32_t rank,
+        std::vector<int64_t> const& sizes,
+        at::ScalarType dtype,
+        int64_t storageOffset)
     {
-        size_t const numel = std::accumulate(sizes.begin(), sizes.end(), 1UL, std::multiplies<size_t>());
-        size_t const elementSize = c10::elementSize(dtype);
-        size_t const reqSize = (numel + storageOffset) * elementSize;
-        TORCH_CHECK(reqSize <= mBufSize, "McastGpuBuffer::getUcBuffer: the requested size (", reqSize,
-            " bytes) exceeds the allocated size (", mBufSize, " bytes)");
-        auto* dataPtr = static_cast<uint8_t*>(mMcastDeviceMemory.getUnicastPtr(rank)) + storageOffset * elementSize;
+        TORCH_CHECK(storageOffset >= 0, "McastGPUBuffer::getUCBuffer: storageOffset must be non-negative");
+        size_t const numel = std::accumulate(
+            sizes.begin(), sizes.end(), size_t{1},
+            [](size_t a, int64_t b) { return a * static_cast<size_t>(b); });
+        size_t const elementSize = c10::elementSize(dtype);
+        size_t const offsetElems = static_cast<size_t>(storageOffset);
+        size_t const reqElems = numel + offsetElems;
+        TORCH_CHECK(
+            elementSize != 0 && reqElems <= mBufSize / elementSize,
+            "McastGPUBuffer::getUCBuffer: requested size (", reqElems * elementSize,
+            " bytes) exceeds allocated size (", mBufSize, " bytes)");
+        auto* basePtr = static_cast<uint8_t*>(mMcastDeviceMemory.getUnicastPtr(rank));
+        auto* dataPtr = basePtr + offsetElems * elementSize;
 
-        auto options = at::TensorOptions().dtype(dtype).device(mLocalDevice);
-        return at::for_blob(dataPtr, c10::IntArrayRef(sizes))
+        auto options = at::TensorOptions().dtype(dtype).device(mLocalDevice);
+        return at::for_blob(dataPtr, c10::IntArrayRef(sizes))
             .options(options)
             .target_device(mLocalDevice)
             .make_tensor();
     }

Also make the header self-contained for these utilities (see include suggestion below).

---

75-89: Mirror the same safety and type fixes in getMCBuffer.

-    at::Tensor getMCBuffer(std::vector<long int> sizes, torch::ScalarType dtype, int64_t storageOffset)
+    at::Tensor getMCBuffer(
+        std::vector<int64_t> const& sizes,
+        at::ScalarType dtype,
+        int64_t storageOffset)
     {
-        size_t const numel = std::accumulate(sizes.begin(), sizes.end(), 1UL, std::multiplies<size_t>());
-        size_t const elementSize = c10::elementSize(dtype);
-        size_t const reqSize = (numel + storageOffset) * elementSize;
-        TORCH_CHECK(reqSize <= mBufSize, "McastGpuBuffer::getMcBuffer: the requested size (", reqSize,
-            " bytes) exceeds the allocated size (", mBufSize, " bytes)");
-        auto* dataPtr = static_cast<uint8_t*>(mMcastDeviceMemory.getMulticastPtr()) + storageOffset * elementSize;
+        TORCH_CHECK(storageOffset >= 0, "McastGPUBuffer::getMCBuffer: storageOffset must be non-negative");
+        size_t const numel = std::accumulate(
+            sizes.begin(), sizes.end(), size_t{1},
+            [](size_t a, int64_t b) { return a * static_cast<size_t>(b); });
+        size_t const elementSize = c10::elementSize(dtype);
+        size_t const offsetElems = static_cast<size_t>(storageOffset);
+        size_t const reqElems = numel + offsetElems;
+        TORCH_CHECK(
+            elementSize != 0 && reqElems <= mBufSize / elementSize,
+            "McastGPUBuffer::getMCBuffer: requested size (", reqElems * elementSize,
+            " bytes) exceeds allocated size (", mBufSize, " bytes)");
+        auto* basePtr = static_cast<uint8_t*>(mMcastDeviceMemory.getMulticastPtr());
+        auto* dataPtr = basePtr + offsetElems * elementSize;
 
-        auto options = at::TensorOptions().dtype(dtype).device(mLocalDevice);
-        return at::for_blob(dataPtr, c10::IntArrayRef(sizes))
+        auto options = at::TensorOptions().dtype(dtype).device(mLocalDevice);
+        return at::for_blob(dataPtr, c10::IntArrayRef(sizes))
             .options(options)
             .target_device(mLocalDevice)
             .make_tensor();
     }