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

fix : correct prefill sampling and record alignment

Browse Source 
- sample the first token from prefill logits without duplicating the prompt tail
- reject incomplete multi-output records before preprocessing alignment
- cover cached generation and partial DPO records with regression tests
This commit is contained in:
ViperEkura
2026-08-05 22:20:29 +08:00
parent a317a4756b
commit 654e6eb0d1
7 changed files with 143 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/core/executor.py
+63 -57
View File
@@ -85,9 +85,65 @@ class Executor:
dtype=self.dtype,
)
def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
def _sample_logits(
self,
logits: Tensor,
tasks: List[Task],
return_logprobs: bool = False,
info: Optional[SamplingBatchInfo] = None,
):
info = info or _build_sampling_batch_info(tasks, self.device)
if info.has_freq:
history_lists = [
t.prompt_ids[-t.rep_window :] + t.output_ids for t in tasks
]
history_lens = [len(ids) for ids in history_lists]
max_len = max(history_lens, default=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, ids in enumerate(history_lists):
length = len(ids)
padded_ids[i, :length] = torch.as_tensor(
ids, dtype=torch.long, device=self.device
)
padded_mask[i, :length] = True
else:
padded_ids = None
padded_mask = None
result = sample(
logits,
temperature=info.temperatures,
top_k=info.top_ks,
top_p=info.top_ps,
frequency_penalty=info.freq_penalties,
input_ids=padded_ids,
input_mask=padded_mask,
return_logprobs=return_logprobs,
)
if not return_logprobs:
return result.tolist()
tokens, logprobs = result
tokens_list = tokens.tolist()
logprobs_list = logprobs.tolist()
for task, logprob in zip(tasks, logprobs_list):
task.output_logprobs.append(float(logprob))
return list(zip(tokens_list, logprobs_list))
def execute_prefill(
self,
tasks: List[Task],
prompt_len: int,
start_pos: int = 0,
return_logprobs: bool = False,
):
if start_pos >= prompt_len:
return
return []
tasks = sorted(tasks, key=lambda t: t.task_id)
batch_sz = len(tasks)
@@ -109,7 +165,7 @@ class Executor:
)
with torch.inference_mode():
self.model(
outputs = self.model(
input_ids,
input_mask=input_mask,
position_ids=position_ids,
@@ -119,6 +175,9 @@ class Executor:
start_pos=start_pos,
),
)
logits = outputs["logits"][:, -1, :]
return tasks, self._sample_logits(logits, tasks, return_logprobs)
def execute_decode(
self, tasks: List[Task], return_logprobs: bool = False
@@ -167,34 +226,6 @@ class Executor:
total_len = max(t.next_pos for t in tasks) + 1
input_mask = self._workspace.decode_mask(position_ids, total_len)
has_freq = info.has_freq
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
else:
padded_ids = None
padded_mask = None
with torch.inference_mode():
outputs = self.model(
input_ids,
@@ -207,29 +238,4 @@ class Executor:
)
logits = outputs["logits"][:, -1, :]
if return_logprobs:
tokens, logprobs = sample(
logits,
temperature=info.temperatures,
top_k=info.top_ks,
top_p=info.top_ps,
frequency_penalty=info.freq_penalties,
input_ids=padded_ids,
input_mask=padded_mask,
return_logprobs=True,
)
tokens_list = tokens.tolist()
logprobs_list = logprobs.tolist()
for t, lp in zip(tasks, logprobs_list):
t.output_logprobs.append(float(lp))
return list(zip(tokens_list, logprobs_list))
return sample(
logits,
temperature=info.temperatures,
top_k=info.top_ks,
top_p=info.top_ps,
frequency_penalty=info.freq_penalties,
input_ids=padded_ids,
input_mask=padded_mask,
).tolist()
return self._sample_logits(logits, tasks, return_logprobs, info=info)
astrai/inference/core/scheduler.py
+21 -15
View File
@@ -91,11 +91,10 @@ class InferenceScheduler:
Single shared primitive for both the continuous-batching loop and
the synchronous ``run_batch`` path, so the two cannot drift.
Tasks must already be allocated in the KV cache. Any task that still
needs prefill (``output_tokens == 0`` and fewer cached pages than
prompt tokens) is prefilled first, grouped by ``(prompt_len, cached)``
so a ragged batch is padded into equal-length groups. Every task is
then extended by one token and decoded.
Tasks must already be allocated in the KV cache. Tasks without output
are prefilled first and sample their first token from the final prompt
position. Tasks with output extend the cache by one position and decode
from their latest generated token.
Args:
tasks: Active tasks to advance by one token.
@@ -109,23 +108,27 @@ class InferenceScheduler:
"""
cache = self._cache
to_prefill = [
t
for t in tasks
if t.output_tokens == 0 and cache.task_cached(t.task_id) < len(t.prompt_ids)
]
to_prefill = [t for t in tasks if t.output_tokens == 0 and t.prompt_ids]
prefilled_ids = set()
produced: List[Task] = []
if to_prefill:
for t in to_prefill:
t.input_tokens = len(t.prompt_ids)
groups: Dict[Tuple[int, int], List[Task]] = {}
for t in to_prefill:
groups.setdefault(
(len(t.prompt_ids), cache.task_cached(t.task_id)), []
).append(t)
start_pos = min(cache.task_cached(t.task_id), len(t.prompt_ids) - 1)
groups.setdefault((len(t.prompt_ids), start_pos), []).append(t)
for (prompt_len, start_pos), group in groups.items():
self._executor.execute_prefill(group, prompt_len, start_pos)
prefilled, step_out = self._executor.execute_prefill(
group, prompt_len, start_pos, return_logprobs=return_logprobs
)
for t, out in zip(prefilled, step_out):
t.output_ids.append(out[0] if return_logprobs else out)
t.output_tokens += 1
prefilled_ids.add(t.task_id)
produced.append(t)
start_logical_page = start_pos // getattr(cache, "page_size", 64)
for t in group:
cache.task_record_hashes(
@@ -135,6 +138,8 @@ class InferenceScheduler:
decoded: List[Task] = []
aborted: List[Task] = []
for t in tasks:
if t.task_id in prefilled_ids:
continue
if cache.task_extend(t.task_id, t.next_pos):
decoded.append(t)
else:
@@ -148,8 +153,9 @@ class InferenceScheduler:
for t, out in zip(decoded, step_out):
t.output_ids.append(out[0] if return_logprobs else out)
t.output_tokens += 1
produced.append(t)
return decoded, aborted
return produced, aborted
def _run_generation_loop(self):
stop_ids = self._task_mgr.tokenizer.stop_ids
astrai/inference/core/task.py
+2 -1
View File
@@ -105,7 +105,8 @@ class Task:
@property
def next_pos(self) -> int:
return self.input_tokens + len(self.output_ids)
# The first output is sampled from prefill and enters KV on the next step.
return self.input_tokens + max(0, len(self.output_ids) - 1)
def is_finished(self, stop_ids: List[int]) -> bool:
if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
astrai/preprocessing/builder.py
+12 -7
View File
@@ -416,7 +416,11 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
return None
result: dict = {}
any_output = False
required_outputs = {
output_key
for output_key, spec in sources_spec.items()
if spec.get("sections")
}
for output_key, spec in sources_spec.items():
sections = spec.get("sections", [])
@@ -428,7 +432,6 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
if ids is None:
continue
result[output_key] = ids
any_output = True
continue
list_field = spec.get("list_field", False)
@@ -444,7 +447,6 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
result[output_key] = ids
if mask is not None:
result[mask_key] = mask
any_output = True
continue
ids, mask = self.renderer.process_sections(
@@ -460,9 +462,7 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
elif "mask_key" in spec:
result[mask_key] = mask
any_output = True
if not any_output:
if not required_outputs or not required_outputs.issubset(result):
return None
result["domain"] = _extract_domain(item, config.output.domain_key)
@@ -474,6 +474,11 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
return [None] * len(items)
results = [{} for _ in items]
required_outputs = {
output_key
for output_key, spec in sources_spec.items()
if spec.get("sections")
}
for output_key, spec in sources_spec.items():
sections = spec.get("sections", [])
if not sections:
@@ -506,7 +511,7 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
return [
({**result, "domain": _extract_domain(item, config.output.domain_key)})
if result
if required_outputs and required_outputs.issubset(result)
else None
for item, result in zip(items, results)
]