Merge pull request #852 from Lander-Hatsune/main

cpuinfer: filter repeated backend instantiation

Merge pull request #852 from Lander-Hatsune/main
cpuinfer: filter repeated backend instantiation
7f57769c · Atream · GitHub · 1c001b80 · d166fb9f · 7f57769c
Unverified Commit 7f57769c authored Mar 14, 2025 by Atream Committed by GitHub Mar 14, 2025
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)