perf: eliminate GPU syncs in contiguous cache write/gather hot paths
- Replace .tolist() calls with _total_len in gather(); move _slot_len updates from per-layer write to once-per-step bind_tasks - Use torch.as_tensor instead of torch.tensor in decode penalty history construction
This commit is contained in:
@@ -435,24 +435,13 @@ class ContiguousCacheView(CacheView):
|
||||
pos = self._write_positions
|
||||
self._cache.k[layer_id, indices, pos] = k.squeeze(1)
|
||||
self._cache.v[layer_id, indices, pos] = v.squeeze(1)
|
||||
for s, p in zip(indices.tolist(), pos.tolist()):
|
||||
cur = self._cache._slot_len.get(s, 0)
|
||||
if p + 1 > cur:
|
||||
self._cache._slot_len[s] = p + 1
|
||||
else:
|
||||
start_pos = self._total_len - seq_len
|
||||
self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k
|
||||
self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v
|
||||
new_len = start_pos + seq_len
|
||||
for s in indices.tolist():
|
||||
cur = self._cache._slot_len.get(s, 0)
|
||||
if new_len > cur:
|
||||
self._cache._slot_len[s] = new_len
|
||||
|
||||
def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
|
||||
max_len = max(
|
||||
self._cache._slot_len.get(int(s), 0) for s in self._batch_indices.tolist()
|
||||
)
|
||||
max_len = self._total_len
|
||||
indices = self._batch_indices
|
||||
k = self._cache.k[layer_id, indices, :max_len]
|
||||
v = self._cache.v[layer_id, indices, :max_len]
|
||||
@@ -528,6 +517,9 @@ class ContiguousCache(KVCache):
|
||||
) -> ContiguousCacheView:
|
||||
slots = [self._task_slot[tid] for tid in task_ids]
|
||||
batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
|
||||
for slot in slots:
|
||||
if total_len > self._slot_len.get(slot, 0):
|
||||
self._slot_len[slot] = total_len
|
||||
return ContiguousCacheView(
|
||||
self, batch_indices, total_len, write_positions=write_positions
|
||||
)
|
||||
|
||||
@@ -104,28 +104,25 @@ class Executor:
|
||||
)
|
||||
|
||||
history_lists = []
|
||||
mask_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)
|
||||
mask_lists.append([True] * len(ids))
|
||||
history_lens.append(len(ids))
|
||||
|
||||
max_len = max(len(h) for h in history_lists)
|
||||
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, m) in enumerate(zip(history_lists, mask_lists)):
|
||||
padded_ids[i, : len(h)] = torch.tensor(
|
||||
h, dtype=torch.long, device=self.device
|
||||
)
|
||||
padded_mask[i, : len(m)] = torch.tensor(
|
||||
m, 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
|
||||
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(
|
||||
|
||||
Reference in New Issue
Block a user