deepinfer-decode
deepinfer/csrc/engine/scheduler.hpp · a679bb8 · L556–659 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
556 /// decode（completion）调度：先把 prompted 队列中的 seq 移入 running，
557 /// 再按 LCFS 挑选本轮 decode 的 seq，block 不足时抢占 output 最短的 seq。
558 /// @note 需要在持有 BlockManagerMutexAcquire 的情况下调用。
559 std::unique_ptr<SchedulerOutputs> _schedule_decode(BlockManagerMutexAcquire &acquire) {
560 (void)acquire; // 仅作为持锁证明，decode 内部直接使用 *_nolock 接口
561
562 std::vector<SequencePtr> finished_seqs = free_finished_seqs();
563
564 // waiting 里已终结的 seq（abort / VE 失败入队）只有调度器能弹出，原本由
565 // _schedule_prefill 顺手清理。现在 prefill 只在本 rank want=true 时才走到，
566 // 某 rank 的 waiting 只剩这类 seq 时 want 恒为 false，它们会永远滞留：
567 // has_unfinished 恒真拖着全组空转，ALL 投票也被这个 rank 永久压死。
568 // 所以在 decode 分支兜底清理；abort 可命中队列中间，需要扫全队列。
569 std::vector<SequencePtr> ignored_seqs;
570 for (auto it = waiting.begin(); it != waiting.end();) {
571 if ((*it)->is_finished()) {
572 ignored_seqs.push_back(std::move(*it));
573 it = waiting.erase(it);
574 } else {
575 ++it;
576 }
577 }
578
579 const size_t max_seqs = max_num_seqs;
580
581 // prompted 队列头部的 seq 依次转入 running（或已结束的转入 finished），
582 // forked seq 需要先克隆 state 和最后一个（可能写了一半的）block。
583 auto admit = [&](SequencePtr &&seq) {
584 if (seq->schedule.status == SequenceStatus::WAITING) {
585 seq->schedule.status = SequenceStatus::RUNNING;
586 allocate_mtp_info_slot(*seq);
587 running.push_back(std::move(seq));
588 } else {
589 TORCH_CHECK(seq->is_finished(), "prompted 队列中的 seq 状态为 ", int(seq->schedule.status), " ，不应该");
590 finished_seqs.push_back(std::move(seq));
591 }
592 prompted.pop_front();
593 };
594
595 while (!prompted.empty() && running.size() < max_seqs) {
596 auto &&[seq, forked] = prompted.front();
597 if (seq->schedule.status == SequenceStatus::FINISHED_ABORTED) { // 进 CE 之前就 abort 了
598 finished_seqs.push_back(std::move(seq));
599 prompted.pop_front();
600 continue;
601 }
602 if (forked) {
603 if ((state_manager && !state_manager->can_allocate_single_nolock()) || !block_manager->can_allocate_single_nolock()) {
604 break;
605 }
606 _clone_state(*seq);
607 int new_block = block_manager->allocate_single_nolock();
608 int num_blocks = (seq->cache.num_tokens + seq->cache.block_size - 1) / seq->cache.block_size;
609 TORCH_CHECK_LE(num_blocks, static_cast<int>(seq->cache.block_table.size()));
610 clone_blocks.emplace_back(seq->cache.block_table[num_blocks - 1]);
611 clone_blocks.emplace_back(new_block);
612 for (size_t i = num_blocks; i < seq->cache.block_table.size(); ++i) {
613 clone_blocks.emplace_back(seq->cache.block_table[i]);
614 clone_blocks.emplace_back(0);
615 }
616 seq->cache.block_table[num_blocks - 1] = new_block;
617 seq->cache.block_table.erase(seq->cache.block_table.begin() + num_blocks, seq->cache.block_table.end());
618 }
619 admit(std::move(seq));
620 }
621
622 // LCFS
623 std::vector<SequencePtr> scheduled;
624 std::sort(running.begin(), running.end(),
625 [](const SequencePtr &a, const SequencePtr &b) { return a->get_output_len() > b->get_output_len(); });
626 scheduled.reserve(max_seqs);
627 size_t l = 0, r = running.size();
628 while (l < r && scheduled.size() < max_seqs) {
629 auto &seq = running[r - 1];
630 while (!block_manager->can_allocate_single_nolock() && r - l > 1) {
631 auto victim_seq = running[l];
632 _preempt(victim_seq, ignored_seqs);
633 l += 1;
634 }
635 if (block_manager->can_allocate_single_nolock()) {
636 _append_slot(*seq);
637 scheduled.push_back(std::move(seq));
638 } else {
639 _preempt(seq, ignored_seqs);
640 }
641
642 r -= 1;
643 }
644
645 TORCH_CHECK(l == r, "decode 时 running 有排队，已经调度 ", scheduled.size(), "个，还剩 ", r - l, " 个，不应该");
646 running.clear();
647
648 std::sort(
649 scheduled.begin(), scheduled.end(), [](const SequencePtr &a, const SequencePtr &b) { return a->get_len() > b->get_len(); });
650
651 std::vector<int> seqlens_q;
652 _mtp_varlen.schedule(scheduled, seqlens_q);
653
654 int num_batched_tokens = 0;
655 for (int seqlen_q : seqlens_q) {
656 num_batched_tokens += seqlen_q;
657 }
658 return _make_outputs(
659 std::move(scheduled), false, num_batched_tokens, std::move(finished_seqs), std::move(ignored_seqs), std::move(seqlens_q), 0);