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

refactor: decode 按页分桶批处理，position_ids 改为 per-task 构建

Browse Source
This commit is contained in:
ViperEkura
2026-05-14 14:22:11 +08:00
parent c0effc9f5b
commit 6269bacfc3
3 changed files with 16 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/scheduler.py
+7 -5
View File
@@ -126,15 +126,17 @@ class InferenceScheduler:
pos_groups: Dict[int, List[Task]] = {}
for t in self._task_mgr.get_active_tasks():
pos_groups.setdefault(t.next_pos, []).append(t)
chunk = t.next_pos // self._page_cache.page_size
key = chunk if chunk <= 1 else 1 << (chunk.bit_length() - 1)
pos_groups.setdefault(key, []).append(t)
if pos_groups:
best_pos = max(pos_groups, key=lambda p: len(pos_groups[p]))
group = sorted(pos_groups[best_pos], key=lambda t: t.task_id)
best_key = max(pos_groups, key=lambda k: len(pos_groups[k]))
group = sorted(pos_groups[best_key], key=lambda t: t.task_id)
valid: List[Task] = []
for t in group:
if self._page_cache.task_extend(t.task_id, best_pos):
if self._page_cache.task_extend(t.task_id, t.next_pos):
valid.append(t)
else:
t.status = TaskStatus.ABORTED
@@ -142,7 +144,7 @@ class InferenceScheduler:
t.stream_callback(STOP)
if valid:
next_tokens = self._executor.execute_decode(valid, best_pos)
next_tokens = self._executor.execute_decode(valid)
for t, ntok in zip(valid, next_tokens):
t.output_ids.append(ntok)