Fix flush cache (#627)

python/sglang/bench_latency.py

@@ -284,23 +284,26 @@ def main(server_args, bench_args):
     else:
         work_func = latency_test
 
-    workers = []
-    for tp_rank in range(server_args.tp_size):
-        proc = multiprocessing.Process(
-            target=work_func,
-            args=(
-                server_args,
-                bench_args,
-                tp_rank,
-            ),
-        )
-        proc.start()
-        workers.append(proc)
+    if server_args.tp_size == 1:
+        work_func(server_args, bench_args, 0)
+    else:
+        workers = []
+        for tp_rank in range(server_args.tp_size):
+            proc = multiprocessing.Process(
+                target=work_func,
+                args=(
+                    server_args,
+                    bench_args,
+                    tp_rank,
+                ),
+            )
+            proc.start()
+            workers.append(proc)
 
-    for proc in workers:
-        proc.join()
+        for proc in workers:
+            proc.join()
 
-    proc.terminate()
+        proc.terminate()
 
 
 if __name__ == "__main__":

python/sglang/srt/managers/controller/manager_single.py

@@ -96,6 +96,7 @@ class ControllerSingle:
     def __init__(self, server_args: ServerArgs, port_args: PortArgs, model_overide_args: dict):
         # Parse args
         self.server_args = server_args
+        self.tp_procs = []
 
         # Init communication
         context = zmq.Context(2)

python/sglang/srt/memory_pool.py

@@ -98,6 +98,8 @@ class TokenToKVPool:
         self.can_use_mem_size += len(free_index)
 
     def clear(self):
+        self.prefetch_buffer = torch.empty(0, device="cuda", dtype=torch.int32)
+
         self.mem_state.fill_(True)
         self.can_use_mem_size = self.size