perf: reduce decode overhead in scheduler and executor
- Precompute page_table and decode_mask on KVCache once per step in PagePool.bind_tasks, instead of per-layer in CudaBackend/TorchNativeBackend - Skip frequency penalty history tensor construction when all penalties are 0 in Executor.execute_decode - Omit FrequencyPenaltyStrategy from sampling pipeline when penalty is 0 - Deduplicate get_active_tasks calls in scheduler loop (3 to 1), remove redundant sorted() on decode tasks - Benchmark (L20, bf16, CUDA backend): B=1 9.48->9.40ms (+1%), B=4 10.73->9.89ms (+8.6%), B=8 10.77->10.13ms (+6.4%)
This commit is contained in:
@@ -203,6 +203,10 @@ class KVCache:
|
||||
seq_lens: [batch_size] — per-request total sequence lengths
|
||||
out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
|
||||
max_len: max(seq_lens) as Python int — avoids GPU sync in decode
|
||||
page_table: [batch, max_len] — precomputed gather indices for decode;
|
||||
None for prefill or when not yet computed.
|
||||
decode_mask: [batch, max_len] bool — precomputed position validity
|
||||
mask for decode; None for prefill or single-batch decode.
|
||||
"""
|
||||
|
||||
k_buffer: Tensor
|
||||
@@ -212,6 +216,8 @@ class KVCache:
|
||||
seq_lens: Tensor
|
||||
out_cache_loc: Tensor
|
||||
max_len: int = 0
|
||||
page_table: Optional[Tensor] = None
|
||||
decode_mask: Optional[Tensor] = None
|
||||
|
||||
|
||||
class PagePool:
|
||||
@@ -428,11 +434,21 @@ class PagePool:
|
||||
out_cache_loc = self._req_pool.req_to_token[
|
||||
req_pool_indices, start_pos:seq_len
|
||||
]
|
||||
page_table = None
|
||||
decode_mask = None
|
||||
else:
|
||||
write_pos = seq_lens_t - 1
|
||||
out_cache_loc = self._req_pool.req_to_token[
|
||||
req_pool_indices, write_pos
|
||||
].unsqueeze(-1)
|
||||
ml = max(seq_lens)
|
||||
page_table = self._req_pool.req_to_token[req_pool_indices, :ml]
|
||||
if len(task_ids) > 1:
|
||||
decode_mask = (
|
||||
torch.arange(ml, device=device)[None, :] < seq_lens_t[:, None]
|
||||
)
|
||||
else:
|
||||
decode_mask = None
|
||||
|
||||
return KVCache(
|
||||
k_buffer=self._storage.k_buffer,
|
||||
@@ -442,6 +458,8 @@ class PagePool:
|
||||
seq_lens=seq_lens_t,
|
||||
out_cache_loc=out_cache_loc,
|
||||
max_len=max(seq_lens),
|
||||
page_table=page_table,
|
||||
decode_mask=decode_mask,
|
||||
)
|
||||
|
||||
# ---- internals ----
|
||||
|
||||
@@ -105,26 +105,33 @@ class Executor:
|
||||
[t.frequency_penalty for t in tasks], device=self.device
|
||||
)
|
||||
|
||||
history_lists = []
|
||||
history_lens = []
|
||||
for t in tasks:
|
||||
window = t.rep_window
|
||||
prompt_part = t.prompt_ids[-window:]
|
||||
ids = prompt_part + t.output_ids
|
||||
history_lists.append(ids)
|
||||
history_lens.append(len(ids))
|
||||
has_freq = bool((freq_penalties != 0).any())
|
||||
if has_freq:
|
||||
history_lists = []
|
||||
history_lens = []
|
||||
for t in tasks:
|
||||
window = t.rep_window
|
||||
prompt_part = t.prompt_ids[-window:]
|
||||
ids = prompt_part + t.output_ids
|
||||
history_lists.append(ids)
|
||||
history_lens.append(len(ids))
|
||||
|
||||
max_len = max(history_lens) if history_lens else 0
|
||||
padded_ids = torch.zeros(
|
||||
len(tasks), max_len, dtype=torch.long, device=self.device
|
||||
)
|
||||
padded_mask = torch.zeros(
|
||||
len(tasks), max_len, dtype=torch.bool, device=self.device
|
||||
)
|
||||
for i, h in enumerate(history_lists):
|
||||
L = history_lens[i]
|
||||
padded_ids[i, :L] = torch.as_tensor(h, dtype=torch.long, device=self.device)
|
||||
padded_mask[i, :L] = True
|
||||
max_len = max(history_lens) if history_lens else 0
|
||||
padded_ids = torch.zeros(
|
||||
len(tasks), max_len, dtype=torch.long, device=self.device
|
||||
)
|
||||
padded_mask = torch.zeros(
|
||||
len(tasks), max_len, dtype=torch.bool, device=self.device
|
||||
)
|
||||
for i, h in enumerate(history_lists):
|
||||
L = history_lens[i]
|
||||
padded_ids[i, :L] = torch.as_tensor(
|
||||
h, dtype=torch.long, device=self.device
|
||||
)
|
||||
padded_mask[i, :L] = True
|
||||
else:
|
||||
padded_ids = None
|
||||
padded_mask = None
|
||||
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(
|
||||
|
||||
@@ -109,9 +109,11 @@ class InferenceScheduler:
|
||||
self._task_mgr.wait_for_tasks(timeout=1.0)
|
||||
continue
|
||||
|
||||
active = self._task_mgr.get_active_tasks()
|
||||
|
||||
to_prefill = [
|
||||
t
|
||||
for t in self._task_mgr.get_active_tasks()
|
||||
for t in active
|
||||
if t.output_tokens == 0
|
||||
and cache.task_cached(t.task_id) < len(t.prompt_ids)
|
||||
]
|
||||
@@ -137,10 +139,10 @@ class InferenceScheduler:
|
||||
t.task_id, t.prompt_ids, start_logical_page
|
||||
)
|
||||
|
||||
decode_tasks = self._task_mgr.get_active_tasks()
|
||||
decode_tasks = active
|
||||
|
||||
valid: List[Task] = []
|
||||
for t in sorted(decode_tasks, key=lambda t: t.task_id):
|
||||
for t in decode_tasks:
|
||||
if cache.task_extend(t.task_id, t.next_pos):
|
||||
valid.append(t)
|
||||
else:
|
||||
|
||||
Reference in New Issue
Block a user