cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp (1)

173-206:Layer-wise send to all connections may deadlock with DP-filtered recv.

unformat() receives only frompickRecvConnections(...); here the layer-wise path always sends on every connection index, which can block when the peer won’trecv some indices under DP gating.

• Restrict sends to a DP-aligned subset (mirrorpickRecvConnections mapping for the destination).
• Or introduce apickSendConnections(...) helper and use it in both layer-wise and non-layer-wise paths.

cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp (2)

163-179:Send only to DP-aligned subset; current code risks over-sending and stalls.

pickRecvConnections(...) narrows receivers toPP-domain entries for the chosen TP slice, butformat() sends on every connection index. Limit sends to the same subset derived fromdestDPRank to prevent unmatched sends.

@@-    auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx);+    auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx);@@-    size_t const pPDomainSize = targetInfo.mDomainPPSize;+    size_t const pPDomainSize = targetInfo.mDomainPPSize;     TLLM_CHECK((cacheBlockSize * blockNum) % pPDomainSize == 0);     auto const targetBufferSize = (cacheBlockSize * blockNum) / pPDomainSize;@@     auto* agentConnnecion = dynamic_cast<executor::kv_cache::AgentConnection const*>(connections[0]);@@-    // The size of outputSplitCaches should be equal to pPDomainSize+    // The size of outputSplitCaches should be equal to pPDomainSize+    // Compute sender indices aligned with the destination DP rank: [tpOffset .. tpOffset+ppSize)+    std::vector<size_t> sendIndices;+    {+        int destDpRank = destConfig.getParallelConfig().mEnableAttentionDP ? destConfig.getParallelConfig().mDPrank : 0;+        size_t const tpOffset = (destDpRank % targetInfo.mDomainTPSize) * targetInfo.mDomainPPSize;+        sendIndices.resize(pPDomainSize);+        std::iota(sendIndices.begin(), sendIndices.end(), tpOffset);+    }@@-    if (connections.size() > 1)+    if (sendIndices.size() > 1)     {         if (!common::getEnvEnableReceiveKVCacheParallel())         {             TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache.");-            for (size_t i = 0; i < connections.size(); i++)+            for (size_t i = 0; i < sendIndices.size(); i++)             {-                sendBufferFun(deviceId, i);+                sendBufferFun(deviceId, sendIndices[i]);             }         }         else         {             // concurrency num-            auto concurrencyNum = std::min(std::max(static_cast<size_t>(1), bufferCoverTargetNum), pPDomainSize);+            auto concurrencyNum = std::min(std::max(static_cast<size_t>(1), bufferCoverTargetNum), sendIndices.size());-            auto remainSendNum = connections.size();+            auto remainSendNum = sendIndices.size();              while (remainSendNum > 0)             {                 auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum);                 std::vector<std::future<void>> futures;                 futures.reserve(sendConcurrencyNum);-                for (size_t i = 0; i < sendConcurrencyNum; i++)+                for (size_t i = 0; i < sendConcurrencyNum; i++)                 {-                    TLLM_CHECK((i + (connections.size() - remainSendNum)) < connections.size());-                    futures.push_back(std::async(-                        std::launch::async, sendBufferFun, deviceId, i + (connections.size() - remainSendNum)));+                    size_t idx = i + (sendIndices.size() - remainSendNum);+                    TLLM_CHECK(idx < sendIndices.size());+                    futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, sendIndices[idx]));                 }                 for (auto& future : futures)                 {                     future.get();                 }                 remainSendNum -= sendConcurrencyNum;             }         }     }     else     {-        sendBufferFun(deviceId, 0);+        TLLM_CHECK(!sendIndices.empty());+        sendBufferFun(deviceId, sendIndices[0]);     }

Also applies to: 250-289

---

563-568:Bug: comparesselfConfig twice; should comparedestConfig for MLA head count.

This always evaluates the same side and misses invalid dest configs.

-    if ((selfConfig.getModelConfig().mNbKvHeadsPerLayer.at(0) != 1)-        || (selfConfig.getModelConfig().mNbKvHeadsPerLayer.at(0) != 1))+    if ((selfConfig.getModelConfig().mNbKvHeadsPerLayer.at(0) != 1)+        || (destConfig.getModelConfig().mNbKvHeadsPerLayer.at(0) != 1))     {         TLLM_LOG_WARNING("MLACacheFormatter::inquireSupport: only support MLA");         return false;     }