deepinfer-schedule-type
deepinfer/csrc/engine/scheduler.hpp · a679bb8 · L275–343 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
275 std::unique_ptr<SchedulerOutputs> _schedule() {
276 py::gil_scoped_release release; // otherwise other python threads would be blocked
277 BlockManagerMutexAcquire acquire{block_manager, state_manager};
278
279 // 上一个 step 的 SchedulerOutputs 借用这三个 vector 的存储（见 pybind 的
280 // clone_states/clone_blocks/zero_states property），Python 侧必须在下一次
281 // schedule 前消费完，这里 clear 后复用。
282 clone_states.clear();
283 zero_states.clear();
284 clone_blocks.clear();
285
286 tokens_chunked_prefill_remaining = 0;
287
288 switch (engine_type) {
289 case EngineType::PROMPT:
290 // PROMPT 引擎只做 prefill，即使没调度到请求也返回（空的）prefill 结果
291 return _schedule_prefill(acquire);
292 case EngineType::COMPLETION:
293 // COMPLETION 引擎只做 decode
294 return _schedule_decode(acquire);
295 case EngineType::NORMAL: {
296 // NORMAL 引擎：决定这一步全组做 prompt 还是 decode。dp 内必须同型，靠投票统一。
297 //
298 // 两类尝试点：
299 // ALL 点（每 prefill_attempt_interval 个 decode step）：全员都有 prefill 活才做，
300 // 防止一个 rank 干活、其余 rank 陪跑空 prompt step。
301 // ANY 点（每 prefill_any_attempt_interval 个 decode step）：一个 rank 有活就一起做，
302 // 给孤立 rank 兜底——它的 prefill 最多等这么多个 decode step。
303 // 两类点重合时按 ANY 算。
304 //
305 // prompt step 会把 decode_step 清零（见 schedule()），所以刚做完 prefill 的下一步
306 // decode_step==0，又落在 ALL 点上：想连着 prefill 同样要全员有活。
307 // decode_step 各 rank 一致，投票的次数和语义全组对得上，不会卡在 allreduce。
308 bool is_any_point = (decode_step != 0) && (decode_step % prefill_any_attempt_interval == 0);
309 bool is_all_point = (decode_step % prefill_attempt_interval == 0);
310 if (!is_any_point && !is_all_point) {
311 return _schedule_decode(acquire);
312 }
313 bool want = (get_gpu_cache_usage() <= GPU_CACHE_PREFILL_WATERMARK) && _has_prefill_waiting();
314 bool do_prefill;
315 if (is_any_point) {
316 do_prefill = (_dp_allreduce(want ? 0 : 1, c10d::ReduceOp::MIN) == 0); // ANY：任一 rank 想做即触发
317 } else {
318 do_prefill = (_dp_allreduce(want ? 1 : 0, c10d::ReduceOp::MIN) == 1); // ALL：全员想做才触发
319 }
320 if (!do_prefill) {
321 return _schedule_decode(acquire);
322 }
323 std::unique_ptr<SchedulerOutputs> outputs;
324 if (want) {
325 outputs = _schedule_prefill(acquire);
326 }
327
328 // 第二轮投票：确认这轮真有 rank 排出了 prefill。want 只看了 waiting 非空，
329 // 但 _schedule_prefill 仍可能空手（slot 满、KV 块不够）。
330 // 全组都空手时必须转 decode：硬做空 prompt step 会把 decode_step 清零，
331 // 下一步接着投票、接着空手——running 满 + waiting 非空就死循环了。
332 bool any_scheduled = (_dp_allreduce(outputs ? 0 : 1, c10d::ReduceOp::MIN) == 0);
333 if (!any_scheduled) {
334 return _schedule_decode(acquire);
335 }
336 if (outputs) {
337 return outputs;
338 }
339 // 组内确有 rank 在做 prompt：本 rank 陪跑空 prompt step 保持步型一致
340 return _make_outputs({}, true, 0, {}, {}, {}, 0);
341 }
342 }
343 TORCH_CHECK(false, "unknown engine_type ", int(engine_type));