[None][chore] optimize kv cache transfer for context TEP and gen DEP #7613

chuangz0 · 2025-09-08T10:31:09Z

Summary by CodeRabbit

New Features
- None
Bug Fixes
- Corrected cache transfer routing in multi-GPU Data Parallel setups to align sends/receives with DP ranks.
- Fixed incorrect “need to send” decisions when DP is disabled, improving correctness across configurations.
Refactor
- Streamlined DP-aware logic and added early exits to avoid unnecessary cache allocation and transfers, reducing overhead.
Tests
- Expanded coverage for DP-enabled and DP-disabled scenarios across MLA and non-MLA cases with updated expectations.

Description

When context TEP and gen DEP, each dp rank of gen will request kv cache from context tp rank0 before.
We changed it so that dp rank will request kv cache from context tp_rank = gen_dp_rank%tp_num.
need TRTLLM_PARALLEL_CACHE_SEND=1

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

chuangz0 · 2025-09-08T10:31:21Z

/bot run

tensorrt-cicd · 2025-09-08T10:37:25Z

PR_Github #18035 [ run ] triggered by Bot

coderabbitai · 2025-09-08T10:40:44Z

📝 Walkthrough

Walkthrough

Adds DP-rank–aware gating for KV-cache send/recv in cache formatting paths, aligning senders and receivers by DupHeadFactor and DP ranks. Updates recipient selection logic, introduces early return when no send is needed, and extends/adjusts unit tests for MLA and non-MLA scenarios with DP enabled/disabled.

Changes

Cohort / File(s)	Summary of changes
DP-aware send/recv gating and routing `cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp`, `cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp`	needSendCache now aligns self/destination DP ranks with DupHeadFactor; adds DP-aware recipient selection in pickRecvConnections; mlaCacheFormatter adds dpRank-based target mapping and an early return in format() when no send is needed; comments updated; no public API changes.
Unit tests: DP scenarios and expectations `cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp`	Adds parameterized DP-enabled MLA tests; adds NoMLA duplicate instantiations; updates expected needSend flags when DP is disabled across several contexts; increases coverage for asymmetric and DP-aware flows.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant Caller
  participant CacheFormatter
  participant Cluster as DP/TP/PP Config
  participant Network

  Caller->>CacheFormatter: format(request)
  CacheFormatter->>Cluster: get self/dest ranks, DupHeadFactor
  CacheFormatter->>CacheFormatter: needSendCache(self, dest, DupHeadFactor)
  alt needSendCache == false
    CacheFormatter-->>Caller: return (no allocation/send)
  else needSendCache == true
    CacheFormatter->>CacheFormatter: pickRecvConnections(selfDP, PeerDupHeadFactor)
    CacheFormatter->>Network: send KV-cache to selected peers
    Network-->>CacheFormatter: ack
    CacheFormatter-->>Caller: return
  end
  note over CacheFormatter,Network: Gating uses (destDPRank % DupHeadFactor) == (selfTpRankInDpGroup % DupHeadFactor).<br/>Recipient selection uses (i % PeerDupHeadFactor) == (selfDPRank % PeerDupHeadFactor).

sequenceDiagram
  autonumber
  participant Caller
  participant MLACacheFormatter as MLA CacheFormatter
  participant Cluster as DP/TP/PP Config
  participant Network

  Caller->>MLACacheFormatter: format(request)
  MLACacheFormatter->>Cluster: read self/dest TP/DP/PP, DupHeadFactor
  MLACacheFormatter->>MLACacheFormatter: needSendCache(...) with DP-aware checks
  alt no send
    MLACacheFormatter-->>Caller: return
  else send
    MLACacheFormatter->>MLACacheFormatter: pickRecvConnections(dpRank-offset mapping)
    MLACacheFormatter->>Network: send KV-cache to DP/PP-aligned targets
    Network-->>MLACacheFormatter: ack
    MLACacheFormatter-->>Caller: return
  end
  note over MLACacheFormatter: targetInfo mapping uses dpRank offset and mDomainTPSize for recipient indices.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

[None][fix]revert kvcache transfer #6709: Touches the same needSendCache/pickRecvConnections functions but removes DP-rank–aware modulo checks; directly intersects logic changed here.
[None][chore] optimize kv cache transfer for context TEP and gen DEP #6657: Also incorporates DP-rank logic into needSendCache and pickRecvConnections in cacheFormatter and mlaCacheFormatter.
[None][feat] move kv cache measure into transfer session #6633: Modifies cacheFormatter/mlaCacheFormatter format/unformat flows overlapping with KV-cache handling paths touched here.

Suggested reviewers

Tabrizian
pcastonguay
Shixiaowei02

✨ Finishing Touches

📝 Generate Docstrings

🧪 Generate unit tests

Create PR with unit tests
Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

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 from pickRecvConnections(...); here the layer-wise path always sends on every connection index, which can block when the peer won’t recv some indices under DP gating.

Restrict sends to a DP-aligned subset (mirror pickRecvConnections mapping for the destination).

Or introduce a pickSendConnections(...) 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 to PP-domain entries for the chosen TP slice, but format() sends on every connection index. Limit sends to the same subset derived from destDPRank 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: compares selfConfig twice; should compare destConfig 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;
     }

🧹 Nitpick comments (3)

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp (1)
131-137: Receiver-subset sanity check.

Nice DP-aware filter. Add an assertion to catch topology mismatches early.
@@
   TLLM_CHECK(numConnections == targetInfo.mIRanks.size());
   int selfDPRank = selfConfig.getParallelConfig().mEnableAttentionDP ? selfConfig.getParallelConfig().mDPrank : 0;
 
   std::vector<size_t> ret;
   for (int i = 0; i < targetInfo.mDomainTPSize; i++)
   {
     if ((i % targetInfo.mPeerDupHeadFactor) == (selfDPRank % targetInfo.mPeerDupHeadFactor))
     {
       for (int j = 0; j < targetInfo.mDomainPPSize; j++)
       {
         ret.push_back((i * targetInfo.mDomainPPSize) + j);
       }
     }
   }
+  TLLM_CHECK(ret.size() * targetInfo.mPeerDupHeadFactor == numConnections);
   return ret;
cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp (1)
1-1: Set required env for parallel send path in tests.

PR requires TRTLLM_PARALLEL_CACHE_SEND=1. Consider setting it once in the test fixture setup to avoid accidental regressions on CI.

Example (add to AsymmetricalCacheTest::setUpCacheTransceiver() before creating connection manager):
@@
     void setUpCacheTransceiver()
     {
+        setenv("TRTLLM_PARALLEL_CACHE_SEND", "1", 1);
         if (!(mIsContext || mIsGeneration))
         {
             return;
         }
cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp (1)
48-54: Recv-indexing aligned to DP rank; add sanity checks.

Add a bound/assert to ensure we get exactly PP-domain connections and avoid surprises if topology changes.
@@
-    // targetInfo , mRanks [tpranks, ppranks]
+    // targetInfo.mIRanks is laid out as [tpRanks, ppRanks]
     int dpRank = selfConfig.getParallelConfig().mEnableAttentionDP ? selfConfig.getParallelConfig().mDPrank : 0;
 
     for (int i = 0; i < targetInfo.mDomainPPSize; i++)
     {
         ret.push_back(i + (dpRank % (targetInfo.mDomainTPSize)) * targetInfo.mDomainPPSize);
     }
+    TLLM_CHECK(ret.size() == static_cast<size_t>(targetInfo.mDomainPPSize));
     return ret;

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between dd9627d and 64b71e2.

📒 Files selected for processing (3)

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp (2 hunks)
cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp (3 hunks)
cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp (3 hunks)

🧰 Additional context used

📓 Path-based instructions (5)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh}