fix(engine): resume preempted sequences by recompute

[nosyndicate]

Fixes recompute-based resume for preempted sequences, which was previously broken end-to-end: every prefill raised ValueError, and (once that's fixed) resumed sequences would get corrupted logits and clobbered metrics.

• Scheduler: ScheduledBatch now carries resumed_sequence_ids, populated in schedule() at the moment a waiting sequence's prior state (PREEMPTED) is still known — before it gets overwritten to RUNNING. Also fixes _preempt to recompute num_tokens as prompt + generated length, since resume re-feeds both.
• _prepare_prefill: uses resumed_sequence_ids (not seq.state, which the scheduler already flips to RUNNING before the engine sees it) to decide whether to feed only prompt tokens (fresh) or prompt + generated tokens (resumed recompute).
• _post_prefill: slices the flattened logits per sequence by its actual fed length (P fresh, P+G resumed) instead of always num_prompt_tokens, fixing misalignment in batches mixing fresh and resumed sequences. Also skips on_prefill_started/num_prompt_tokens reassignment on resume, since resume is a second prefill for an already-admitted request and would otherwise clobber start_ns (corrupting queue_wait_ms/ttft_ms/total_ms).
• Adds regression tests: resumed-sequence token feeding, mixed fresh+resumed batch logit slicing, start_ns/metrics preservation across resume, and an end-to-end forced-preemption test asserting byte-identical output vs. an uncontended solo run.

[Copilot]

Pull request overview

Fixes recompute-based resume for preempted sequences by explicitly tracking which sequences are resumed at scheduling time, and then using that information in the engine to (a) feed the correct tokens during resume-prefill and (b) slice flattened logits correctly in mixed fresh+resumed batches without corrupting metrics.

Changes:

• Scheduler: recomputes num_tokens on preemption and threads resumed_sequence_ids through ScheduledBatch for PREFILL batches.
• Engine: uses resumed_sequence_ids to feed prompt vs prompt+generated, and to slice prefill logits by per-sequence fed length while avoiding metric clobbering on resume.
• Tests: adds regression coverage for token feeding, mixed-batch logit slicing, start/metrics preservation, and an end-to-end forced-preemption equivalence test.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/executor/test_scheduler.py	Adds regression tests for preemption num_tokens recompute and for resumed_sequence_ids being carried in scheduled PREFILL batches.
tests/executor/test_schedule_engine.py	Adds unit + end-to-end tests for resume-by-recompute behavior (token feeding, logits slicing, metrics preservation, forced-preemption output equivalence).
server/executor/types.py	Extends ScheduledBatch to include resumed_sequence_ids (default empty).
server/executor/scheduler.py	Marks resumed sequences during PREFILL scheduling and recomputes num_tokens at preemption time.
server/executor/engine.py	Uses scheduler-provided resumed_sequence_ids to control resume-prefill inputs, per-seq logits slicing, and to avoid resetting metrics on resume.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.