perf: 消除非流式推理 CPU 空转并减少 decode GPU 张量冗余分配

- engine.py: _Result 改用 threading.Condition.wait_for 替代
  Event busy-wait,非流式模式线程被内核挂起而非 1760 万次空转
- scheduler.py: _execute_decode 将 temperature/top_k/top_p 张量
  移至循环外预先分配,避免每步重复 torch.tensor();input_ids
  改用 torch.empty 避免不必要的 zero 初始化(两处均为完全覆盖)
- _execute_prefill: input_ids 同改为 torch.empty
This commit is contained in:
2026-05-10 15:32:11 +08:00
parent 3583c46b66
commit cffedaad5e
2 changed files with 31 additions and 17 deletions
+18 -8
View File
@@ -97,7 +97,8 @@ class _Result:
"""Thread-safe token accumulator for streaming and non-streaming modes.
Supports multiple concurrent generation tasks with per-index result tracking.
Uses a threading.Event for efficient waiting on completion.
Uses a threading.Condition for efficient completion notification
and a threading.Event for streaming wakeup.
"""
def __init__(self, count: int = 1):
@@ -106,7 +107,7 @@ class _Result:
Args:
count: Number of concurrent generation tasks to track.
"""
self._lock = threading.Lock()
self._cond = threading.Condition()
self._event = threading.Event()
self.tokens: List[str] = []
self.results: List[str] = [""] * count
@@ -124,7 +125,7 @@ class _Result:
token: The decoded token string, or STOP sentinel.
idx: Index of the generation task this token belongs to.
"""
with self._lock:
with self._cond:
self.tokens.append(token)
if token is not STOP:
self.results[idx] += token
@@ -132,7 +133,8 @@ class _Result:
if not self._done[idx]:
self._done[idx] = True
self._completed += 1
self._event.set()
self._cond.notify_all()
self._event.set()
def pop_all(self) -> List[str]:
"""Returns and clears all accumulated tokens.
@@ -140,7 +142,7 @@ class _Result:
Returns:
List of token strings since the last call.
"""
with self._lock:
with self._cond:
out = self.tokens.copy()
self.tokens.clear()
if not out:
@@ -158,13 +160,22 @@ class _Result:
"""
return self._event.wait(timeout=timeout)
def wait_completion(self) -> None:
"""Blocks until all tasks complete (non-streaming).
Uses a Condition to sleep efficiently instead of busy-waiting.
The calling thread is parked until a STOP signal arrives.
"""
with self._cond:
self._cond.wait_for(lambda: self._completed >= self._total)
def get_results(self) -> List[str]:
"""Returns all accumulated results for non-streaming mode.
Returns:
List of complete generated strings, one per task index.
"""
with self._lock:
with self._cond:
return self.results.copy()
@@ -425,8 +436,7 @@ class InferenceEngine:
)
task_ids.append(task_id)
while result._completed < result._total:
result.wait(timeout=1.0)
result.wait_completion()
for task_id in task_ids:
self.scheduler.remove_task(task_id)