[None][chore] optimize kv cache transfer for context TEP and gen DEP (#6657)

chuangz0 · web-flow · commit ee471df07c16 · 2025-08-07T11:36:05.000+08:00
Signed-off-by: Chuang Zhu &lt;111838961+chuangz0@users.noreply.github.com&gt;
diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp
@@ -75,7 +75,6 @@ BlockRange getBlockRangeForReceiving(BaseKVCacheManager* cacheManager, LlmReques
 bool CacheFormatter::needSendCache(
     CacheState const& selfConfig, CacheState const& destConfig, runtime::SizeType32 selfIdx)
 {
-    // int selfTpRank = selfIdx % selfConfig.getParallelConfig().mTensorParallelism;
     auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx);
     if (targetInfo.mDupHeadFactor <= 1)
     {
@@ -90,8 +89,9 @@ bool CacheFormatter::needSendCache(
             = selfConfig.getParallelConfig().mTensorParallelism / selfConfig.getParallelConfig().mDPsize;
         selfTpRankInDpGroup = selfTpRank % selfTPNumInDPGroup;
     }
+    int destDPRank = destConfig.getParallelConfig().mEnableAttentionDP ? destConfig.getParallelConfig().mDPrank : 0;
 
-    return selfTpRankInDpGroup % targetInfo.mDupHeadFactor == 0;
+    return (destDPRank % targetInfo.mDupHeadFactor) == (selfTpRankInDpGroup % targetInfo.mDupHeadFactor);
 }
 
 void checkAlternateWindow(BaseKVCacheManager* cacheManager, BaseCacheFormatter::CacheState const& selfConfig,
@@ -128,11 +128,12 @@ std::vector<size_t> CacheFormatter::pickRecvConnections(
         return ret;
     }
     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 == 0)
+        if ((i % targetInfo.mPeerDupHeadFactor) == (selfDPRank % targetInfo.mPeerDupHeadFactor))
         {
             for (int j = 0; j < targetInfo.mDomainPPSize; j++)
             {
diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp
@@ -45,10 +45,12 @@ std::vector<size_t> MLACacheFormatter::pickRecvConnections(
     auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx);
     TLLM_CHECK(numConnections == targetInfo.mIRanks.size());
     std::vector<size_t> ret;
-    // targetInfo , mRanks [tpranks, dpranks]
+    // targetInfo , mRanks [tpranks, ppranks]
+    int dpRank = selfConfig.getParallelConfig().mEnableAttentionDP ? selfConfig.getParallelConfig().mDPrank : 0;
+
     for (int i = 0; i < targetInfo.mDomainPPSize; i++)
     {
-        ret.push_back(i);
+        ret.push_back(i + (dpRank % (targetInfo.mDomainTPSize)) * targetInfo.mDomainPPSize);
     }
     return ret;
 }
@@ -58,19 +60,24 @@ bool MLACacheFormatter::needSendCache(
 {
     int selfTpRank = selfIdx % selfConfig.getParallelConfig().mTensorParallelism;
 
+    int destTPNumInDPGroup = destConfig.getParallelConfig().mEnableAttentionDP
+        ? destConfig.getParallelConfig().mTensorParallelism / destConfig.getParallelConfig().mDPsize
+        : destConfig.getParallelConfig().mTensorParallelism;
+    int destDPRank = destConfig.getParallelConfig().mEnableAttentionDP ? destConfig.getParallelConfig().mDPrank : 0;
+
     if (selfConfig.getParallelConfig().mEnableAttentionDP)
     {
         int selfTPNumInDPGroup
             = selfConfig.getParallelConfig().mTensorParallelism / selfConfig.getParallelConfig().mDPsize;
-        int destTPNumInDPGroup = destConfig.getParallelConfig().mEnableAttentionDP
-            ? destConfig.getParallelConfig().mTensorParallelism / destConfig.getParallelConfig().mDPsize
-            : destConfig.getParallelConfig().mTensorParallelism;
+
         int selfTPrankINDPGroup = selfTpRank % selfTPNumInDPGroup;
         if (selfTPNumInDPGroup <= destTPNumInDPGroup)
         {
             return true;
         }
-        return selfTPrankINDPGroup % (selfTPNumInDPGroup / destTPNumInDPGroup) == 0;
+
+        int dupHeadFactor = selfTPNumInDPGroup / destTPNumInDPGroup;
+        return selfTPrankINDPGroup % dupHeadFactor == destDPRank;
     }
 
     int destTPNum = destConfig.getParallelConfig().mEnableAttentionDP
@@ -81,7 +88,8 @@ bool MLACacheFormatter::needSendCache(
     {
         return true;
     }
-    return selfTpRank % (selfTPNum / destTPNum) == 0;
+    int dupHeadFactor = selfTPNum / destTPNum;
+    return selfTpRank % dupHeadFactor == destDPRank;
 }
 
 void MLACacheFormatter::format(TransferSession& session)
diff --git a/cpp/tests/batch_manager/cacheTransceiverTest.cpp b/cpp/tests/batch_manager/cacheTransceiverTest.cpp
@@ -1457,12 +1457,15 @@ TEST(targetTest, CacheStateNODP)
 
     verifyContext(
         /*contextRank*/ 0, /*expectRanks*/ {0}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1, /*expectNeedSend*/ true);
+
     verifyContext(
         /*contextRank*/ 1, /*expectRanks*/ {0}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1, /*expectNeedSend*/ false);
+
     verifyContext(
         /*contextRank*/ 2, /*expectRanks*/ {1}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1, /*expectNeedSend*/ true);
     verifyContext(
         /*contextRank*/ 3, /*expectRanks*/ {1}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1, /*expectNeedSend*/ false);
+
     verifyContext(
         /*contextRank*/ 4, /*expectRanks*/ {2}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1, /*expectNeedSend*/ true);
     verifyContext(
@@ -1474,7 +1477,6 @@ TEST(targetTest, CacheStateNODP)
 
     contextTP = 2;
     genTP = 4;
-
     verifyContext(
         /*contextRank*/ 0, /*expectRanks*/ {0, 1}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 2, /*expectNeedSend*/ true);
     verifyContext(/*contextRank*/ 1, /*expectRanks*/ {2, 3}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 2,
@@ -1564,13 +1566,13 @@ TEST(targetTest, CacheStateContextDP)
         /*expectNeedSend*/ true);
     verifyContext(
         /*contextRank*/ 0, /*generationRank*/ 1, /*expectRanks*/ {1}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1,
-        /*expectNeedSend*/ true);
+        /*expectNeedSend*/ false);
     verifyContext(
         /*contextRank*/ 1, /*generationRank*/ 0, /*expectRanks*/ {0}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1,
         /*expectNeedSend*/ false);
     verifyContext(
         /*contextRank*/ 1, /*generationRank*/ 1, /*expectRanks*/ {1}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1,
-        /*expectNeedSend*/ false);
+        /*expectNeedSend*/ true);
     verifyContext(
         /*contextRank*/ 2, /*generationRank*/ 0, /*expectRanks*/ {0}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1,
         /*expectNeedSend*/ false);

Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,6 @@ BlockRange getBlockRangeForReceiving(BaseKVCacheManager* cacheManager, LlmReques`
`75`	`75`	`bool CacheFormatter::needSendCache(`
`76`	`76`	`CacheState const& selfConfig, CacheState const& destConfig, runtime::SizeType32 selfIdx)`
`77`	`77`	`{`
`78`		`- // int selfTpRank = selfIdx % selfConfig.getParallelConfig().mTensorParallelism;`
`79`	`78`	`auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx);`
`80`	`79`	`if (targetInfo.mDupHeadFactor <= 1)`
`81`	`80`	`{`
`@@ -90,8 +89,9 @@ bool CacheFormatter::needSendCache(`
`90`	`89`	`= selfConfig.getParallelConfig().mTensorParallelism / selfConfig.getParallelConfig().mDPsize;`
`91`	`90`	`selfTpRankInDpGroup = selfTpRank % selfTPNumInDPGroup;`
`92`	`91`	`}`
	`92`	`+ int destDPRank = destConfig.getParallelConfig().mEnableAttentionDP ? destConfig.getParallelConfig().mDPrank : 0;`
`93`	`93`
`94`		`- return selfTpRankInDpGroup % targetInfo.mDupHeadFactor == 0;`
	`94`	`+ return (destDPRank % targetInfo.mDupHeadFactor) == (selfTpRankInDpGroup % targetInfo.mDupHeadFactor);`
`95`	`95`	`}`
`96`	`96`
`97`	`97`	`void checkAlternateWindow(BaseKVCacheManager* cacheManager, BaseCacheFormatter::CacheState const& selfConfig,`
`@@ -128,11 +128,12 @@ std::vector<size_t> CacheFormatter::pickRecvConnections(`
`128`	`128`	`return ret;`
`129`	`129`	`}`
`130`	`130`	`TLLM_CHECK(numConnections == targetInfo.mIRanks.size());`
	`131`	`+ int selfDPRank = selfConfig.getParallelConfig().mEnableAttentionDP ? selfConfig.getParallelConfig().mDPrank : 0;`
`131`	`132`
`132`	`133`	`std::vector<size_t> ret;`
`133`	`134`	`for (int i = 0; i < targetInfo.mDomainTPSize; i++)`
`134`	`135`	`{`
`135`		`- if (i % targetInfo.mPeerDupHeadFactor == 0)`
	`136`	`+ if ((i % targetInfo.mPeerDupHeadFactor) == (selfDPRank % targetInfo.mPeerDupHeadFactor))`
`136`	`137`	`{`
`137`	`138`	`for (int j = 0; j < targetInfo.mDomainPPSize; j++)`
`138`	`139`	`{`
Original file line number	Diff line number	Diff line change
`@@ -45,10 +45,12 @@ std::vector<size_t> MLACacheFormatter::pickRecvConnections(`
`45`	`45`	`auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx);`
`46`	`46`	`TLLM_CHECK(numConnections == targetInfo.mIRanks.size());`
`47`	`47`	`std::vector<size_t> ret;`
`48`		`- // targetInfo , mRanks [tpranks, dpranks]`
	`48`	`+ // targetInfo , mRanks [tpranks, ppranks]`
	`49`	`+ int dpRank = selfConfig.getParallelConfig().mEnableAttentionDP ? selfConfig.getParallelConfig().mDPrank : 0;`
	`50`	`+`
`49`	`51`	`for (int i = 0; i < targetInfo.mDomainPPSize; i++)`
`50`	`52`	`{`
`51`		`- ret.push_back(i);`
	`53`	`+ ret.push_back(i + (dpRank % (targetInfo.mDomainTPSize)) * targetInfo.mDomainPPSize);`
`52`	`54`	`}`
`53`	`55`	`return ret;`
`54`	`56`	`}`
`@@ -58,19 +60,24 @@ bool MLACacheFormatter::needSendCache(`
`58`	`60`	`{`
`59`	`61`	`int selfTpRank = selfIdx % selfConfig.getParallelConfig().mTensorParallelism;`
`60`	`62`
	`63`	`+ int destTPNumInDPGroup = destConfig.getParallelConfig().mEnableAttentionDP`
	`64`	`+ ? destConfig.getParallelConfig().mTensorParallelism / destConfig.getParallelConfig().mDPsize`
	`65`	`+ : destConfig.getParallelConfig().mTensorParallelism;`
	`66`	`+ int destDPRank = destConfig.getParallelConfig().mEnableAttentionDP ? destConfig.getParallelConfig().mDPrank : 0;`
	`67`	`+`
`61`	`68`	`if (selfConfig.getParallelConfig().mEnableAttentionDP)`
`62`	`69`	`{`
`63`	`70`	`int selfTPNumInDPGroup`
`64`	`71`	`= selfConfig.getParallelConfig().mTensorParallelism / selfConfig.getParallelConfig().mDPsize;`
`65`		`- int destTPNumInDPGroup = destConfig.getParallelConfig().mEnableAttentionDP`
`66`		`- ? destConfig.getParallelConfig().mTensorParallelism / destConfig.getParallelConfig().mDPsize`
`67`		`- : destConfig.getParallelConfig().mTensorParallelism;`
	`72`	`+`
`68`	`73`	`int selfTPrankINDPGroup = selfTpRank % selfTPNumInDPGroup;`
`69`	`74`	`if (selfTPNumInDPGroup <= destTPNumInDPGroup)`
`70`	`75`	`{`
`71`	`76`	`return true;`
`72`	`77`	`}`
`73`		`- return selfTPrankINDPGroup % (selfTPNumInDPGroup / destTPNumInDPGroup) == 0;`
	`78`	`+`
	`79`	`+ int dupHeadFactor = selfTPNumInDPGroup / destTPNumInDPGroup;`
	`80`	`+ return selfTPrankINDPGroup % dupHeadFactor == destDPRank;`
`74`	`81`	`}`
`75`	`82`
`76`	`83`	`int destTPNum = destConfig.getParallelConfig().mEnableAttentionDP`
`@@ -81,7 +88,8 @@ bool MLACacheFormatter::needSendCache(`
`81`	`88`	`{`
`82`	`89`	`return true;`
`83`	`90`	`}`
`84`		`- return selfTpRank % (selfTPNum / destTPNum) == 0;`
	`91`	`+ int dupHeadFactor = selfTPNum / destTPNum;`
	`92`	`+ return selfTpRank % dupHeadFactor == destDPRank;`
`85`	`93`	`}`
`86`	`94`
`87`	`95`	`void MLACacheFormatter::format(TransferSession& session)`