fix: reset cache-aware counters on prefill registration

lightllm/server/httpserver_for_pd_master/manager.py

@@ -664,6 +664,12 @@ def register_pd(self, pd_info_json, websocket):
         if pd_client.mode == "prefill":
             self.prefill_nodes = [e for e in self.prefill_nodes if e.client_ip_port != pd_client.client_ip_port]
             self.prefill_nodes.append(pd_client)
+            # dispatched_prompt_chars is the cumulative counter used by the
+            # CacheAware policy to balance request dispatch across prefill nodes.
+            # Reset all counters together when a node registers or reconnects so
+            # stale history does not bias CacheAware toward the zero-valued node.
+            for prefill_node in self.prefill_nodes:
+                prefill_node.dispatched_prompt_chars = 0
         elif pd_client.mode == "decode":
             self.decode_nodes = [e for e in self.decode_nodes if e.client_ip_port != pd_client.client_ip_port]
             self.decode_nodes.append(pd_client)

unit_tests/server/test_pd_master_mode.py

@@ -184,6 +184,60 @@ def test_pd_manager_without_connected_nodes_is_healthy():
     assert asyncio.run(manager.check_pd_nodes_health()) is True
 
 
+def test_prefill_registration_resets_all_dispatched_prompt_chars():
+    args = StartArgs()
+    manager = PDManager(args)
+
+    def register_prefill(node_id, client_ip_port):
+        manager.register_pd(
+            {
+                "node_id": node_id,
+                "client_ip_port": client_ip_port,
+                "mode": "prefill",
+                "start_args": {
+                    "max_req_total_len": args.max_req_total_len,
+                    "max_image_pixels": args.max_image_pixels,
+                    "disable_image_resize": args.disable_image_resize,
+                },
+            },
+            websocket=object(),
+        )
+
+    register_prefill(1, "10.0.0.1:8000")
+    manager.prefill_nodes[0].dispatched_prompt_chars = 1234
+
+    register_prefill(2, "10.0.0.2:8000")
+
+    assert [node.dispatched_prompt_chars for node in manager.prefill_nodes] == [0, 0]
+
+
+def test_prefill_reconnection_resets_all_dispatched_prompt_chars():
+    args = StartArgs()
+    manager = PDManager(args)
+
+    def pd_info(node_id, client_ip_port):
+        return {
+            "node_id": node_id,
+            "client_ip_port": client_ip_port,
+            "mode": "prefill",
+            "start_args": {
+                "max_req_total_len": args.max_req_total_len,
+                "max_image_pixels": args.max_image_pixels,
+                "disable_image_resize": args.disable_image_resize,
+            },
+        }
+
+    manager.register_pd(pd_info(1, "10.0.0.1:8000"), websocket=object())
+    manager.register_pd(pd_info(2, "10.0.0.2:8000"), websocket=object())
+    manager.prefill_nodes[0].dispatched_prompt_chars = 100
+    manager.prefill_nodes[1].dispatched_prompt_chars = 200
+
+    manager.register_pd(pd_info(3, "10.0.0.2:8000"), websocket=object())
+
+    assert [node.client_ip_port for node in manager.prefill_nodes] == ["10.0.0.1:8000", "10.0.0.2:8000"]
+    assert [node.dispatched_prompt_chars for node in manager.prefill_nodes] == [0, 0]
+
+
 def test_pd_master_inference_health_matches_normal_node_semantics(monkeypatch):
     monkeypatch.setattr("lightllm.server.httpserver_for_pd_master.manager.time.time", lambda: 1000)
     manager = SimpleNamespace(