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

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