diff --git a/csrc/engine/pybind_scheduler.cpp b/csrc/engine/pybind_scheduler.cpp index a379e85..62fabe1 100644 --- a/csrc/engine/pybind_scheduler.cpp +++ b/csrc/engine/pybind_scheduler.cpp @@ -91,6 +91,10 @@ void pybind_engine_scheduler(pybind11::module &engine) { }, py::arg("seq")) .def("add_sequence", &Scheduler::add_sequence, py::arg("seq")) + .def("free_finished_seqs", &Scheduler::free_finished_seqs, + "把 running 中已结束的 seq 移出并返回(不释放 block/state,调用方按需 free_seq);PROMPT 引擎送走 prefill 结果后用") + .def("add_prompted_sequence", &Scheduler::add_prompted_sequence, py::arg("seq"), + "PD 分离:把 KV/state 已填好的 seq 放入 prompted 队列,由 COMPLETION 引擎直接续写") .def("abort_sequence", &Scheduler::abort_sequence, py::arg("request_ids")) .def("has_unfinished_seqs", &Scheduler::has_unfinished_seqs) .def("get_num_unfinished_seqs", &Scheduler::get_num_unfinished_seqs) diff --git a/csrc/engine/scheduler.hpp b/csrc/engine/scheduler.hpp index ed6e897..88e1109 100644 --- a/csrc/engine/scheduler.hpp +++ b/csrc/engine/scheduler.hpp @@ -1058,6 +1058,19 @@ struct Scheduler { } } + /// PD 分离:把在 prefill 引擎上算完 prefill、KV block 与 state 已由调用方填好的 seq 直接放入 prompted + /// 队列(不 clone)。token 为 prompt + 首个生成 token,其中除最后一个 token 外 KV 均已就绪。 + void add_prompted_sequence(std::shared_ptr seq) { + TORCH_CHECK(seq->schedule.status == SequenceStatus::WAITING, "status ", int(seq->schedule.status)); + TORCH_CHECK(!seq->cache.block_table.empty(), "prompted sequence needs allocated KV blocks"); + TORCH_CHECK(seq->cache.num_tokens == seq->get_len() - 1, "prompted sequence must have all but its last token computed: ", + seq->cache.num_tokens, " vs len ", seq->get_len()); + seq->request.arrival_time = get_timestamp_ms(); + seq->reserve_output_tokens(); + seq->cache.block_table.reserve(_block_table_capacity(*seq)); + prompted.emplace_back(std::move(seq), false); + } + SequencePtr fork_seq(const SequencePtr &seq, py::object blocks, std::string &&target_branch_id, py::object sampling_params, py::object so_matcher, int sp_token, int replace_last_token_id, int64_t seq_id, double arrival_time) { TORCH_CHECK(seq->cache.num_tokens > 0, "Cannot fork a sequence without a computed prefix"); diff --git a/deepinfer/_C/engine.pyi b/deepinfer/_C/engine.pyi index f500ad7..08b5a78 100644 --- a/deepinfer/_C/engine.pyi +++ b/deepinfer/_C/engine.pyi @@ -313,6 +313,10 @@ class Scheduler: ... def add_sequence(self, seq: Sequence) -> None: ... + def add_prompted_sequence(self, seq: Sequence) -> None: + ... + def free_finished_seqs(self) -> list[Sequence]: + ... def fork_seq(self, seq: Sequence, blocks: typing.Any, target_branch_id: str, sampling_params: typing.Any, so_matcher: typing.Any, sp_token: typing.SupportsInt | typing.SupportsIndex, replace_last_token_id: typing.SupportsInt | typing.SupportsIndex, seq_id: typing.SupportsInt | typing.SupportsIndex, arrival_time: typing.SupportsFloat | typing.SupportsIndex) -> Sequence: ... def free_seq(self, seq: Sequence) -> None: