revert kvcache transfer

chuangz0 · chuangz0 · commit bcc7a76a5c9a · 2025-08-08T07:53:01.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,6 +75,7 @@ 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)
     {
@@ -89,9 +90,8 @@ bool CacheFormatter::needSendCache(
             = selfConfig.getParallelConfig().mTensorParallelism / selfConfig.getParallelConfig().mDPsize;
         selfTpRankInDpGroup = selfTpRank % selfTPNumInDPGroup;
     }
-    int destDPRank = destConfig.getParallelConfig().mEnableAttentionDP ? destConfig.getParallelConfig().mDPrank : 0;
 
-    return (destDPRank % targetInfo.mDupHeadFactor) == (selfTpRankInDpGroup % targetInfo.mDupHeadFactor);
+    return selfTpRankInDpGroup % targetInfo.mDupHeadFactor == 0;
 }
 
 void checkAlternateWindow(BaseKVCacheManager* cacheManager, BaseCacheFormatter::CacheState const& selfConfig,
@@ -128,12 +128,11 @@ 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) == (selfDPRank % targetInfo.mPeerDupHeadFactor))
+        if (i % targetInfo.mPeerDupHeadFactor == 0)
         {
             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,12 +45,10 @@ 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, ppranks]
-    int dpRank = selfConfig.getParallelConfig().mEnableAttentionDP ? selfConfig.getParallelConfig().mDPrank : 0;
-
+    // targetInfo , mRanks [tpranks, dpranks]
     for (int i = 0; i < targetInfo.mDomainPPSize; i++)
     {
-        ret.push_back(i + (dpRank % (targetInfo.mDomainTPSize)) * targetInfo.mDomainPPSize);
+        ret.push_back(i);
     }
     return ret;
 }
@@ -60,24 +58,19 @@ 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;
         }
-
-        int dupHeadFactor = selfTPNumInDPGroup / destTPNumInDPGroup;
-        return selfTPrankINDPGroup % dupHeadFactor == destDPRank;
+        return selfTPrankINDPGroup % (selfTPNumInDPGroup / destTPNumInDPGroup) == 0;
     }
 
     int destTPNum = destConfig.getParallelConfig().mEnableAttentionDP
@@ -88,8 +81,7 @@ bool MLACacheFormatter::needSendCache(
     {
         return true;
     }
-    int dupHeadFactor = selfTPNum / destTPNum;
-    return selfTpRank % dupHeadFactor == destDPRank;
+    return selfTpRank % (selfTPNum / destTPNum) == 0;
 }
 
 void MLACacheFormatter::format(TransferSession& session)
diff --git a/cpp/tests/batch_manager/cacheTransceiverTest.cpp b/cpp/tests/batch_manager/cacheTransceiverTest.cpp
@@ -1457,15 +1457,12 @@ 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(
@@ -1477,6 +1474,7 @@ 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,
@@ -1566,13 +1564,13 @@ TEST(targetTest, CacheStateContextDP)
         /*expectNeedSend*/ true);
     verifyContext(
         /*contextRank*/ 0, /*generationRank*/ 1, /*expectRanks*/ {1}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1,
-        /*expectNeedSend*/ false);
+        /*expectNeedSend*/ true);
     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*/ true);
+        /*expectNeedSend*/ false);
     verifyContext(
         /*contextRank*/ 2, /*generationRank*/ 0, /*expectRanks*/ {0}, /*expectPPDomain*/ 1, /*expectTPDomain*/ 1,
         /*expectNeedSend*/ false);

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