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:
+        return self._sample_logits(logits, tasks, return_logprobs, info=info)
-            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()

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
+        Tasks must already be allocated in the KV cache. Tasks without output
-        needs prefill (``output_tokens == 0`` and fewer cached pages than
+        are prefilled first and sample their first token from the final prompt
-        prompt tokens) is prefilled first, grouped by ``(prompt_len, cached)``
+        position. Tasks with output extend the cache by one position and decode
-        so a ragged batch is padded into equal-length groups. Every task is
+        from their latest generated token.
-        then extended by one token and decoded.
 
         Args:
             tasks: Active tasks to advance by one token.
@@ -109,23 +108,27 @@ class InferenceScheduler:
         """
         cache = self._cache
 
-        to_prefill = [
+        to_prefill = [t for t in tasks if t.output_tokens == 0 and t.prompt_ids]
-            t
+        prefilled_ids = set()
-            for t in tasks
+        produced: List[Task] = []
-            if t.output_tokens == 0 and cache.task_cached(t.task_id) < len(t.prompt_ids)
-        ]
         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(
+                start_pos = min(cache.task_cached(t.task_id), len(t.prompt_ids) - 1)
-                    (len(t.prompt_ids), cache.task_cached(t.task_id)), []
+                groups.setdefault((len(t.prompt_ids), start_pos), []).append(t)
-                ).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:

astrai/preprocessing/builder.py

@@ -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 required_outputs or not required_outputs.issubset(result):
-
-        if not any_output:
             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)
         ]

tests/data/test_preprocess_builder.py

@@ -369,6 +369,19 @@ def test_dpo_missing_field_is_none(chat_tokenizer, builder):
     assert builder.build({"chosen": [], "rejected": []}, config, chat_tokenizer) is None
 
 
+@pytest.mark.parametrize("missing", ["chosen", "rejected"])
+def test_dpo_partial_record_is_none(chat_tokenizer, builder, missing):
+    config = make_dpo_chat_config()
+    item = {
+        "chosen": [{"role": "assistant", "content": "Good"}],
+        "rejected": [{"role": "assistant", "content": "Bad"}],
+    }
+    item.pop(missing)
+
+    assert builder.build(item, config, chat_tokenizer) is None
+    assert builder.build_batch([item], config, chat_tokenizer) == [None]
+
+
 def test_grpo_basic(chat_tokenizer, builder):
     config = make_grpo_config()
     item = {

tests/inference/test_scheduler.py

@@ -205,6 +205,36 @@ def test_run_batch_returns_token_sequences(device):
         scheduler.stop()
 
 
+def test_run_batch_tokens_match_full_sequence_forward(device):
+    scheduler, _tok, model = _make_real_scheduler(device)
+    prompt = [10, 20, 30, 40]
+    try:
+        expected = []
+        sequence = list(prompt)
+        for _ in range(2):
+            input_ids = torch.tensor([sequence], dtype=torch.long, device=device)
+            position_ids = torch.arange(len(sequence), device=device).unsqueeze(0)
+            input_mask = torch.ones(
+                1, len(sequence), len(sequence), dtype=torch.bool, device=device
+            ).tril()
+            with torch.inference_mode():
+                logits = model(
+                    input_ids,
+                    input_mask=input_mask,
+                    position_ids=position_ids,
+                )["logits"][:, -1, :]
+            token = logits.argmax(dim=-1).item()
+            expected.append(token)
+            sequence.append(token)
+
+        result = scheduler.run_batch(
+            prompt_ids_list=[prompt], max_tokens=2, temperature=0
+        )
+        assert result == [expected]
+    finally:
+        scheduler.stop()
+
+
 def test_run_batch_return_logprobs_aligned(device):
     """return_logprobs=True gives (token_ids, logprobs) tuples with equal len."""
     scheduler, _tok, _model = _make_real_scheduler(device)

tests/inference/test_task.py

@@ -22,6 +22,8 @@ def test_task_next_pos():
     task.input_tokens = 5
     assert task.next_pos == 5
     task.output_ids.append(4)
+    assert task.next_pos == 5
+    task.output_ids.append(5)
     assert task.next_pos == 6