cpuinfer: filter repeated backend instantiation

d166fb9f · Lander-Hatsune · 09c043d8 · d166fb9f
Commit d166fb9f authored Mar 10, 2025 by Lander-Hatsune
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

ktransformers/operators/cpuinfer.py ktransformers/operators/cpuinfer.py +6 -1

No files found.
--- a/ktransformers/operators/cpuinfer.py
+++ b/ktransformers/operators/cpuinfer.py
@@ -727,8 +727,13 @@ class CPUInferKVCache:

 class CPUInfer:
    cpuinfer = None
+    cur_backend_thread_num = 0
+    
    def __init__(self, thread_num):
-        CPUInfer.cpuinfer = cpuinfer_ext.CPUInfer(thread_num)
+        if thread_num > CPUInfer.cur_backend_thread_num:
+            CPUInfer.cur_backend_thread_num = thread_num
+            del CPUInfer.cpuinfer
+            CPUInfer.cpuinfer = cpuinfer_ext.CPUInfer(thread_num)

    def submit(self, task):
        CPUInfer.cpuinfer.submit(task)