fix kv_transfer_stat reduce bug

juncgu-google · juncgu-google · commit 81d8fc000b3d · 2025-11-22T02:29:56.000Z
Signed-off-by: Juncheng Gu &lt;jcgu@google.com&gt;
diff --git a/tpu_inference/distributed/offload/tpu_offload_connector.py b/tpu_inference/distributed/offload/tpu_offload_connector.py
@@ -312,8 +312,8 @@ def reduce(self) -> dict[str, int | float]:
                 "Num finished load blocks ": 0,
             }
 
-        finished_save_chunks = sum(self.data["finished_save_chunks"].values())
-        finished_load_chunks = sum(self.data["finished_load_chunks"].values())
+        finished_save_chunks = sum(len(chunk_list) for chunk_list in self.data["finished_save_chunks"].values())
+        finished_load_chunks = sum(len(chunk_list) for chunk_list in self.data["finished_load_chunks"].values())
 
         return {
             "Num finished save chunks ": finished_save_chunks,

Original file line number	Diff line number	Diff line change
`@@ -312,8 +312,8 @@ def reduce(self) -> dict[str, int \| float]:`
`312`	`312`	`"Num finished load blocks ": 0,`
`313`	`313`	`}`
`314`	`314`
`315`		`- finished_save_chunks = sum(self.data["finished_save_chunks"].values())`
`316`		`- finished_load_chunks = sum(self.data["finished_load_chunks"].values())`
	`315`	`+ finished_save_chunks = sum(len(chunk_list) for chunk_list in self.data["finished_save_chunks"].values())`
	`316`	`+ finished_load_chunks = sum(len(chunk_list) for chunk_list in self.data["finished_load_chunks"].values())`
`317`	`317`
`318`	`318`	`return {`
`319`	`319`	`"Num finished save chunks ": finished_save_chunks,`