[TRTLLM-13969][feat] Support MiniMax M3 for Disaggregated Serving

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -216,12 +216,20 @@ struct VmmDescSplitter
     /// For non-VRAM or addresses not in the map, descs pass through unchanged.
     [[nodiscard]] static MemoryDescs splitDescsWithRegionMap(MemoryDescs const& descs, VramRegionMap const& regionMap);
 
-    /// @brief Split paired src/dst descs using local and remote region maps.
-    /// src is split by localRegionMap, dst is split by remoteRegionMap.
-    /// The final piece size is min(srcPiece, dstPiece, remaining).
-    [[nodiscard]] static std::pair<MemoryDescs, MemoryDescs> splitTransferDescsWithRegionMaps(
-        MemoryDescs const& srcDescs, MemoryDescs const& dstDescs, VramRegionMap const& localRegionMap,
-        VramRegionMap const& remoteRegionMap);
+    /// @brief Split paired src/dst descs at chunk boundaries, then coalesce contiguous pieces.
+    /// src is split by localRegionMap, dst is split by remoteRegionMap; each piece size is
+    /// min(srcPiece, dstPiece, remaining). Pairs are sorted by src address, and adjacent pieces
+    /// whose src AND dst are both contiguous (same deviceId) are merged — but a merged desc never
+    /// crosses a chunk boundary on either side, and never spans two distinct regions, so every
+    /// output desc stays within a single registered memory region. Merging requires region
+    /// metadata: a piece whose address misses the region map on either side is never merged,
+    /// because two unknown regions are indistinguishable and a merge could cross a chunk or
+    /// registration boundary. With no region metadata the result is split-only. Non-kVRAM descs
+    /// pass through unchanged (no region info is available to bound the merge).
+    /// @param enableCoalesce When false, only split at chunk boundaries without merging pieces.
+    [[nodiscard]] static std::pair<MemoryDescs, MemoryDescs> splitAndCoalesceTransferDescs(MemoryDescs const& srcDescs,
+        MemoryDescs const& dstDescs, VramRegionMap const& localRegionMap, VramRegionMap const& remoteRegionMap,
+        bool enableCoalesce = true);
 
     /// @brief Split VRAM descs at VMM chunk boundaries detected via cuMemGetAddressRange.
     /// For cudaMalloc memory (single allocation), descs pass through unchanged.

cpp/tensorrt_llm/common/envUtils.cpp

@@ -517,10 +517,10 @@ uint16_t getEnvNixlPort()
     return nixlPort;
 }
 
-bool getEnvNixlEnableCoalesce()
+bool getEnvNixlDisableCoalesce()
 {
-    static bool const enableCoalesce = getBoolEnv("TRTLLM_NIXL_ENABLE_COALESCE");
-    return enableCoalesce;
+    static bool const disableCoalesce = getBoolEnv("TRTLLM_NIXL_DISABLE_COALESCE");
+    return disableCoalesce;
 }
 
 bool getEnvDisaggBenchmarkGenOnly()

cpp/tensorrt_llm/common/envUtils.h

@@ -151,7 +151,8 @@ bool getEnvKVCachePoolUseFabricMemory();
 
 uint16_t getEnvNixlPort();
 
-bool getEnvNixlEnableCoalesce();
+// Whether to disable coalescing of contiguous NIXL transfer descriptors (coalescing is on by default).
+bool getEnvNixlDisableCoalesce();
 
 bool getEnvDisaggBenchmarkGenOnly();
 

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp

@@ -347,164 +347,6 @@ NixlTransferStatus::~NixlTransferStatus() noexcept
     }
 }
 
-[[nodiscard]] MemoryDescs NixlHelper::coalesceMemoryDescs(MemoryDescs const& descs)
-{
-    auto const& descVec = descs.getDescs();
-
-    // If empty or single element, return as-is
-    if (descVec.size() <= 1)
-    {
-        return descs;
-    }
-
-    size_t const numDescs = descVec.size();
-
-    // Create index array and sort by address
-    std::vector<size_t> sortedIndices(numDescs);
-    std::iota(sortedIndices.begin(), sortedIndices.end(), 0);
-
-    std::sort(sortedIndices.begin(), sortedIndices.end(),
-        [&descVec](size_t lhs, size_t rhs)
-        {
-            // Sort by deviceId first, then by address
-            if (descVec[lhs].getDeviceId() != descVec[rhs].getDeviceId())
-            {
-                return descVec[lhs].getDeviceId() < descVec[rhs].getDeviceId();
-            }
-            return descVec[lhs].getAddr() < descVec[rhs].getAddr();
-        });
-
-    std::vector<MemoryDesc> coalesced;
-    coalesced.reserve(numDescs);
-
-    // Start with the first entry
-    size_t firstIdx = sortedIndices[0];
-    uintptr_t currentAddr = descVec[firstIdx].getAddr();
-    size_t currentLen = descVec[firstIdx].getLen();
-    uint32_t currentDeviceId = descVec[firstIdx].getDeviceId();
-
-    for (size_t idx = 1; idx < numDescs; ++idx)
-    {
-        size_t sortedIdx = sortedIndices[idx];
-        auto const& desc = descVec[sortedIdx];
-
-        // Check if current can be coalesced with previous
-        bool isContiguous = (currentAddr + currentLen == desc.getAddr()) && (currentDeviceId == desc.getDeviceId());
-
-        if (isContiguous)
-        {
-            // Coalesce: extend the current region
-            currentLen += desc.getLen();
-        }
-        else
-        {
-            // Cannot coalesce: save the current region and start a new one
-            coalesced.emplace_back(currentAddr, currentLen, currentDeviceId);
-
-            currentAddr = desc.getAddr();
-            currentLen = desc.getLen();
-            currentDeviceId = desc.getDeviceId();
-        }
-    }
-
-    // Add the last region
-    coalesced.emplace_back(currentAddr, currentLen, currentDeviceId);
-
-    TLLM_LOG_DEBUG("NixlHelper::coalesceMemoryDescs: coalesced %zu -> %zu entries", descVec.size(), coalesced.size());
-
-    return MemoryDescs{descs.getType(), std::move(coalesced)};
-}
-
-[[nodiscard]] std::pair<MemoryDescs, MemoryDescs> NixlHelper::coalesceTransferDescs(
-    TransferDescs const& srcDescs, TransferDescs const& dstDescs)
-{
-    auto const& srcVec = srcDescs.getDescs();
-    auto const& dstVec = dstDescs.getDescs();
-
-    // If sizes don't match or empty, return as-is
-    if (srcVec.size() != dstVec.size() || srcVec.empty())
-    {
-        return {srcDescs, dstDescs};
-    }
-
-    size_t const numDescs = srcVec.size();
-
-    // Create index array and sort by src address
-    // This allows us to find contiguous regions even if the original order is scattered
-    std::vector<size_t> sortedIndices(numDescs);
-    std::iota(sortedIndices.begin(), sortedIndices.end(), 0);
-
-    std::sort(sortedIndices.begin(), sortedIndices.end(),
-        [&srcVec](size_t lhs, size_t rhs)
-        {
-            // Sort by deviceId first, then by address
-            if (srcVec[lhs].getDeviceId() != srcVec[rhs].getDeviceId())
-            {
-                return srcVec[lhs].getDeviceId() < srcVec[rhs].getDeviceId();
-            }
-            return srcVec[lhs].getAddr() < srcVec[rhs].getAddr();
-        });
-
-    std::vector<MemoryDesc> coalescedSrc;
-    std::vector<MemoryDesc> coalescedDst;
-    coalescedSrc.reserve(numDescs);
-    coalescedDst.reserve(numDescs);
-
-    // Start with the first entry (using sorted order)
-    size_t firstIdx = sortedIndices[0];
-    uintptr_t currentSrcAddr = srcVec[firstIdx].getAddr();
-    size_t currentSrcLen = srcVec[firstIdx].getLen();
-    uint32_t currentSrcDeviceId = srcVec[firstIdx].getDeviceId();
-
-    uintptr_t currentDstAddr = dstVec[firstIdx].getAddr();
-    size_t currentDstLen = dstVec[firstIdx].getLen();
-    uint32_t currentDstDeviceId = dstVec[firstIdx].getDeviceId();
-
-    for (size_t idx = 1; idx < numDescs; ++idx)
-    {
-        size_t sortedIdx = sortedIndices[idx];
-        auto const& src = srcVec[sortedIdx];
-        auto const& dst = dstVec[sortedIdx];
-
-        // Check if current src and dst can be coalesced with previous
-        bool srcContiguous
-            = (currentSrcAddr + currentSrcLen == src.getAddr()) && (currentSrcDeviceId == src.getDeviceId());
-        bool dstContiguous
-            = (currentDstAddr + currentDstLen == dst.getAddr()) && (currentDstDeviceId == dst.getDeviceId());
-
-        if (srcContiguous && dstContiguous)
-        {
-            // Coalesce: extend the current region
-            currentSrcLen += src.getLen();
-            currentDstLen += dst.getLen();
-        }
-        else
-        {
-            // Cannot coalesce: save the current region and start a new one
-            coalescedSrc.emplace_back(currentSrcAddr, currentSrcLen, currentSrcDeviceId);
-            coalescedDst.emplace_back(currentDstAddr, currentDstLen, currentDstDeviceId);
-
-            currentSrcAddr = src.getAddr();
-            currentSrcLen = src.getLen();
-            currentSrcDeviceId = src.getDeviceId();
-
-            currentDstAddr = dst.getAddr();
-            currentDstLen = dst.getLen();
-            currentDstDeviceId = dst.getDeviceId();
-        }
-    }
-
-    // Don't forget to add the last region
-    coalescedSrc.emplace_back(currentSrcAddr, currentSrcLen, currentSrcDeviceId);
-    coalescedDst.emplace_back(currentDstAddr, currentDstLen, currentDstDeviceId);
-
-    TLLM_LOG_DEBUG(
-        "NixlHelper::coalesceTransferDescs: coalesced %zu -> %zu transfer entries", srcVec.size(), coalescedSrc.size());
-
-    return {MemoryDescs{srcDescs.getType(), std::move(coalescedSrc)},
-        MemoryDescs{dstDescs.getType(), std::move(coalescedDst)}};
-}
-
 TransferState NixlTransferStatus::wait(int64_t timeout_ms) const
 {
     auto startTime = std::chrono::steady_clock::now();
@@ -693,12 +535,8 @@ void NixlTransferAgent::registerMemory(RegisterDescs const& descs)
     auto detectedRegionMap = VmmDescSplitter::detectVramRegionMap(descs);
     mLocalVramRegionInfo.merge(detectedRegionMap);
 
-    // Coalesce contiguous memory regions to reduce registration overhead (disabled by default)
-    // Set TRTLLM_NIXL_ENABLE_COALESCE=1 to enable this optimization
-    auto coalescedDescs = common::getEnvNixlEnableCoalesce() ? NixlHelper::coalesceMemoryDescs(splitDescs) : splitDescs;
-
     nixl_status_t status;
-    status = mRawAgent->registerMem(NixlHelper::convertRegDlist(coalescedDescs), &mExtraParams);
+    status = mRawAgent->registerMem(NixlHelper::convertRegDlist(splitDescs), &mExtraParams);
     TLLM_CHECK(status == NIXL_SUCCESS);
 
     std::string localMD;
@@ -713,12 +551,8 @@ void NixlTransferAgent::deregisterMemory(RegisterDescs const& descs)
     // Split using per-region registry info to match what was registered
     auto splitDescs = VmmDescSplitter::splitDescsWithRegionMap(descs, mLocalVramRegionInfo);
 
-    // Coalesce contiguous memory regions to match what was registered (disabled by default)
-    // Set TRTLLM_NIXL_ENABLE_COALESCE=1 to enable this optimization
-    auto coalescedDescs = common::getEnvNixlEnableCoalesce() ? NixlHelper::coalesceMemoryDescs(splitDescs) : splitDescs;
-
     nixl_status_t status;
-    status = mRawAgent->deregisterMem(NixlHelper::convertRegDlist(coalescedDescs), &mExtraParams);
+    status = mRawAgent->deregisterMem(NixlHelper::convertRegDlist(splitDescs), &mExtraParams);
     TLLM_CHECK(status == NIXL_SUCCESS);
 
     // Remove entries from registry
@@ -743,7 +577,7 @@ void NixlTransferAgent::loadRemoteAgent(std::string const& name, AgentDesc const
         name == remoteName, "loadRemoteAgent gets error agent name: %s != %s", name.c_str(), remoteName.c_str());
 
     // Store remote VMM region info for chunk boundary calculations in
-    // VmmDescSplitter::splitTransferDescsWithRegionMaps. Per-agent map because different remote agents may have
+    // VmmDescSplitter::splitAndCoalesceTransferDescs. Per-agent map because different remote agents may have
     // overlapping virtual addresses.
     auto const& regions = agentDesc.getVramRegions();
     if (!regions.empty())
@@ -764,14 +598,13 @@ AgentDesc NixlTransferAgent::getLocalAgentDesc()
     nixl_status_t status = mRawAgent->getLocalMD(nixlBlob);
     TLLM_CHECK(status == NIXL_SUCCESS);
 
-    // Pack local VMM region info so remote agents can compute chunk boundaries.
+    // Pack ALL local region info (VMM multi-chunk and single-allocation alike) so remote agents can
+    // compute chunk boundaries and never coalesce transfer descs across separately registered regions.
     std::vector<VramRegionMeta> regions;
+    regions.reserve(mLocalVramRegionInfo.size());
     for (auto const& [base, info] : mLocalVramRegionInfo)
     {
-        if (info.chunkSize > 0)
-        {
-            regions.push_back({base, info.totalLen, info.chunkSize});
-        }
+        regions.push_back({base, info.totalLen, info.chunkSize});
     }
 
     return AgentDesc{nixlBlob, std::move(regions)};
@@ -809,32 +642,25 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name)
     {
         reqParams.hasNotif = false;
     }
-    // Split transfer descriptors at VMM chunk boundaries to match registered memory.
-    // Both src and dst are split at chunk boundaries to ensure each descriptor
-    // falls within a single registered memory region on both local and remote sides.
-    // Find remote agent's VMM region map (empty map if not found).
+    // Split transfer descriptors at VMM chunk boundaries to match registered memory, then coalesce
+    // contiguous pieces. A coalesced descriptor never crosses a chunk boundary or a registered
+    // region boundary on either side, so every descriptor still falls within a single registered
+    // memory region on both local and remote sides. Set TRTLLM_NIXL_DISABLE_COALESCE=1 to fall back
+    // to split-only descriptors. Find remote agent's region map (empty map if not found — e.g. the
+    // peer's AgentDesc carried no region info; addresses missing from a map are never coalesced,
+    // so an empty remote map degrades to split-only rather than risking merges across unknown
+    // remote chunk/registration boundaries).
     static VramRegionMap const kEmptyMap;
     auto remoteIt = mRemoteVramRegionInfo.find(request.getRemoteName());
     auto const& remoteRegionMap = (remoteIt != mRemoteVramRegionInfo.end()) ? remoteIt->second : kEmptyMap;
 
-    auto [splitSrc, splitDst] = VmmDescSplitter::splitTransferDescsWithRegionMaps(
-        request.getSrcDescs(), request.getDstDescs(), mLocalVramRegionInfo, remoteRegionMap);
+    auto [xferSrc, xferDst] = VmmDescSplitter::splitAndCoalesceTransferDescs(request.getSrcDescs(),
+        request.getDstDescs(), mLocalVramRegionInfo, remoteRegionMap, !common::getEnvNixlDisableCoalesce());
 
-    // Coalesce contiguous memory regions to reduce transfer count (disabled by default)
-    // This matches the coalescing done during registerMemory()
-    // Set TRTLLM_NIXL_ENABLE_COALESCE=1 to enable this optimization
-    if (common::getEnvNixlEnableCoalesce())
-    {
-        NVTX3_SCOPED_RANGE(coalesceTransferDescs_CreateXferReq);
-        auto [coalescedSrc, coalescedDst] = NixlHelper::coalesceTransferDescs(splitSrc, splitDst);
-        status
-            = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(coalescedSrc),
-                NixlHelper::convertXferDist(coalescedDst), request.getRemoteName(), handle, &reqParams);
-    }
-    else
     {
-        status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(splitSrc),
-            NixlHelper::convertXferDist(splitDst), request.getRemoteName(), handle, &reqParams);
+        NVTX3_SCOPED_RANGE(createXferReq);
+        status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(xferSrc),
+            NixlHelper::convertXferDist(xferDst), request.getRemoteName(), handle, &reqParams);
     }
 
     TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS,

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h

@@ -39,21 +39,6 @@ struct NixlHelper
     [[nodiscard]] static nixl_xfer_dlist_t convertXferDist(FileDescs const& descs);
     static void posixGpuToFileFallback(MemoryDescs const& memoryDesc, FileDescs const& fileDescs);
     static void posixFileToGpuFallback(MemoryDescs const& memoryDesc, FileDescs const& fileDescs);
-
-    /// @brief Coalesce contiguous memory regions to reduce memory registration overhead.
-    /// Adjacent memory regions with the same deviceId will be merged into a single region.
-    /// @param descs Memory descriptors to coalesce
-    /// @return Coalesced MemoryDescs
-    [[nodiscard]] static MemoryDescs coalesceMemoryDescs(MemoryDescs const& descs);
-
-    /// @brief Coalesce contiguous memory regions in src and dst to reduce transfer count.
-    /// If src[i] and src[i+1] are contiguous, and dst[i] and dst[i+1] are also contiguous
-    /// (with same deviceId), they will be merged into a single transfer.
-    /// @param srcDescs Source memory descriptors
-    /// @param dstDescs Destination memory descriptors
-    /// @return Pair of coalesced (src, dst) MemoryDescs
-    [[nodiscard]] static std::pair<MemoryDescs, MemoryDescs> coalesceTransferDescs(
-        TransferDescs const& srcDescs, TransferDescs const& dstDescs);
 };
 
 class NixlTransferStatus final : public TransferStatus