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
+13 -9
View File
@@ -253,7 +253,7 @@ class InferenceScheduler:
batch_sz = len(tasks)
seq_len = prompt_len - start_pos
input_ids = torch.zeros(batch_sz, seq_len, dtype=torch.long, device=self.device)
input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
input_mask = torch.ones(batch_sz, seq_len, dtype=torch.bool, device=self.device)
for i, t in enumerate(tasks):
@@ -285,15 +285,21 @@ class InferenceScheduler:
for t in tasks:
self._maybe_alloc_page(t, start_pos)
input_ids = torch.zeros(batch_sz, dtype=torch.long, device=self.device)
for i, t in enumerate(tasks):
input_ids[i] = t.output_ids[-1] if t.output_ids else t.prompt_ids[-1]
input_ids = torch.tensor(
[t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
dtype=torch.long,
device=self.device,
)
active_mask = torch.ones((batch_sz, 1), dtype=torch.bool, device=self.device)
page_tables = self._make_page_table_tensor(tasks)
total_len = start_pos + 1
temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
with torch.inference_mode():
outputs = self.model(
input_ids.unsqueeze(1),
@@ -305,11 +311,9 @@ class InferenceScheduler:
next_tokens = sample(
logits,
temperature=torch.tensor(
[t.temperature for t in tasks], device=logits.device
),
top_k=torch.tensor([t.top_k for t in tasks], device=logits.device),
top_p=torch.tensor([t.top_p for t in tasks], device=logits.device),
temperature=temperatures,
top_k=top_ks,
top_p=top_ps,
).tolist()
for t, ntok in zip(tasks, next_tokens):