ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: merge decode batch for 10x throughput

Browse Source 
- merge all active decode tasks into single forward pass (was grouped by next_pos)
- add per-task write_positions to ContiguousCacheView for correct KV writes
- override ContiguousCache.task_cached (base returned 0, caused prefill loops)
- add --cache_len/--frequency_penalty/--rep_window to generate.py
- chunked batch processing with tqdm progress

bench (1.2B model, 128 prompts, 64 tok, batch=128):
  before: 77.2s, ~111 tok/s
  after:   7.1s, ~1210 tok/s (10.9x)
This commit is contained in:
ViperEkura
2026-07-18 08:50:46 +08:00
parent f7df02f9a3
commit a24a7b4da5
4 changed files with 162 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/core/cache.py
+47 -13
View File
@@ -300,7 +300,11 @@ class KVCache(ABC):
@abstractmethod
def bind_tasks(
self, task_ids: List[str], total_len: int, device: torch.device
self,
task_ids: List[str],
total_len: int,
device: torch.device,
write_positions: Optional[Tensor] = None,
) -> CacheView: ...
def task_cached(self, task_id: str) -> int:
@@ -399,7 +403,11 @@ class PageCache(KVCache):
self._pool.record(page_table[i], prompt_ids, i)
def bind_tasks(
self, task_ids: List[str], total_len: int, device: torch.device
self,
task_ids: List[str],
total_len: int,
device: torch.device,
write_positions: Optional[Tensor] = None,
) -> PageCacheView:
page_table = self._table.table_tensor(task_ids, device)
return PageCacheView(self._storage, page_table, total_len)
@@ -409,23 +417,37 @@ class ContiguousCacheView(CacheView):
"""Contiguous KV-cache view for attention layers."""
def __init__(
self, cache: "ContiguousCache", batch_indices: Tensor, total_len: int = 0
self,
cache: "ContiguousCache",
batch_indices: Tensor,
total_len: int = 0,
write_positions: Optional[Tensor] = None,
):
self._cache = cache
self._batch_indices = batch_indices
self._total_len = total_len
self._write_positions = write_positions
def write(self, layer_id: int, k: Tensor, v: Tensor):
seq_len = k.size(1)
start_pos = self._total_len - seq_len
indices = self._batch_indices
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
if self._write_positions is not None and seq_len == 1:
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(
@@ -491,9 +513,21 @@ class ContiguousCache(KVCache):
def task_extend(self, task_id: str, pos: int) -> bool:
return pos < self.max_seq_len
def task_cached(self, task_id: str) -> int:
slot = self._task_slot.get(task_id)
if slot is None:
return 0
return self._slot_len.get(slot, 0)
def bind_tasks(
self, task_ids: List[str], total_len: int, device: torch.device
self,
task_ids: List[str],
total_len: int,
device: torch.device,
write_positions: Optional[Tensor] = None,
) -> ContiguousCacheView:
slots = [self._task_slot[tid] for tid in task_ids]
batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
return ContiguousCacheView(self, batch_indices, total_len)
return ContiguousCacheView(
self, batch_indices, total_len, write_positions=write_positions
)
astrai/inference/core/executor.py
+6 -1
View File
@@ -106,7 +106,12 @@ class Executor:
with torch.inference_mode():
outputs = self.model(
input_ids.unsqueeze(1),
paged_cache=self.kv_cache.bind_tasks(task_ids, total_len, self.device),
paged_cache=self.kv_cache.bind_tasks(
task_ids,
total_len,
self.device,
write_positions=position_ids,
),
position_ids=position_ids.unsqueeze(1),
)
logits = outputs["logits"][:, -1, :]
astrai/inference/core/scheduler.py
+21 -26
View File
@@ -138,36 +138,31 @@ class InferenceScheduler:
t.task_id, t.prompt_ids, start_logical_page
)
pos_groups: Dict[int, List[Task]] = {}
for t in self._task_mgr.get_active_tasks():
pos_groups.setdefault(t.next_pos, []).append(t)
decode_tasks = self._task_mgr.get_active_tasks()
for next_pos in sorted(pos_groups.keys()):
group = sorted(pos_groups[next_pos], key=lambda t: t.task_id)
valid: List[Task] = []
for t in sorted(decode_tasks, key=lambda t: t.task_id):
if cache.task_extend(t.task_id, t.next_pos):
valid.append(t)
else:
t.status = TaskStatus.ABORTED
self._task_mgr.invoke_callback(t.task_id, STOP)
valid: List[Task] = []
for t in group:
if cache.task_extend(t.task_id, t.next_pos):
valid.append(t)
else:
t.status = TaskStatus.ABORTED
if valid:
next_tokens = self._executor.execute_decode(valid)
for t, ntok in zip(valid, next_tokens):
t.output_ids.append(ntok)
t.output_tokens += 1
self._task_mgr.invoke_callback(
t.task_id,
self._task_mgr.tokenizer.decode([ntok]),
)
for t in valid:
if t.is_finished(stop_ids):
self._task_mgr.invoke_callback(t.task_id, STOP)
if valid:
next_tokens = self._executor.execute_decode(valid)
for t, ntok in zip(valid, next_tokens):
t.output_ids.append(ntok)
t.output_tokens += 1
self._task_mgr.invoke_callback(
t.task_id,
self._task_mgr.tokenizer.decode([ntok]),
)
for t in valid:
if t.is_finished(stop_ids):
self._task_mgr.invoke_callback(t.task_id, STOP)
except Exception as e:
self._stop_event.set()
logger.error(f"Scheduler loop crashed: {e}", exc_info=True)