fix : correct prefill sampling and record alignment

- 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

astrai/inference/core/executor.py

@@ -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

@@ -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

@@ -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: