deepinfer-metadata

deepinfer/csrc/engine/scheduler.hpp · a679bb8 · L667–734 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

667    py::tuple schedule(SequenceMetadataPool *sequence_metadata_pool) {
668        auto p_scheduler_outputs = _schedule();
669        auto &scheduler_outputs = *p_scheduler_outputs;
670
671        TORCH_CHECK(scheduler_outputs.num_batched_tokens <= max_num_batched_tokens,
672            "调度的请求 token 数量超过了 max_num_batched_tokens, 不应该", "max_num_batched_tokens: ", max_num_batched_tokens,
673            " scheduler_outputs.num_batched_tokens: ", scheduler_outputs.num_batched_tokens);
674
675        schedule_counter += 1;
676        if (scheduler_outputs.prompt_run) {
677            decode_step = 0;
678        } else {
679            running = scheduler_outputs.scheduled_seqs;
680            decode_step += 1;
681            if (nextn > 0) {
682                TORCH_CHECK(get_num_used_mtp_info_slots() == static_cast<int64_t>(running.size()),
683                    "The number of used MTP info slots must match the running queue size: used=", get_num_used_mtp_info_slots(),
684                    ", running=", running.size());
685            }
686        }
687
688        py::list seq_metadata_list(scheduler_outputs.scheduled_seqs.size());
689        int max_seqlen_q = nextn + 1;
690        int64_t attn_flops = 0, total_flops = 0;
691        int i = 0;
692        for (auto &&p_seq : scheduler_outputs.scheduled_seqs) {
693            auto &&seq = p_seq.get();
694            seq->mtp.schedule_cnt += 1;
695            auto golden_token = seq->mtp.next_golden_tokens(seq->get_output_len(), nextn);
696            int seqlen_q = scheduler_outputs.seqlens_q[i];
697            int seqlen_k = seq->cache.num_tokens + seq->schedule.total_inflight_tokens + seqlen_q;
698            // 本步 position 的开区间上界：prefill 止于 seqlen_k；decode 的 seqlen_q 在 GPU 上至多补到 max_seqlen_q（同 zos.cpp 的
699            // max_k），DFlash draft 再往后 nextn。RoPE cache 等按 position 索引的 buffer 都只有 runtime_max_model_len 行。
700            int position_end = scheduler_outputs.prompt_run ? seqlen_k : seqlen_k - seqlen_q + max_seqlen_q + nextn;
701            TORCH_CHECK(position_end <= runtime_max_model_len,
702                std::format("position_end={} exceeds runtime_max_model_len={}: request_id={} branch={} len={} num_tokens={} inflight={} "
703                            "seqlen_q={} prompt_run={}",
704                    position_end, runtime_max_model_len, seq->request.request_id, seq->schedule.branch_id, seq->get_len(),
705                    seq->cache.num_tokens, seq->schedule.total_inflight_tokens, seqlen_q, scheduler_outputs.prompt_run));
706            auto &&p_metadata =
707                sequence_metadata_pool->allocate(seq, scheduler_outputs.prompt_run, seqlen_q, seqlen_k, std::move(golden_token));
708
709            scheduler_outputs.sum_seqlen_qxk += static_cast<int64_t>(seqlen_q) * seqlen_k;
710            scheduler_outputs.max_seqlen_k = std::max(scheduler_outputs.max_seqlen_k, seqlen_k);
711
712            if (log_stats) {
713                const bool is_prompt = scheduler_outputs.prompt_run;
714                const int full_len = seq->input.tokens.get_len();
715                const auto [attn, total] =
716                    flops_estimator->calc_flops(seqlen_q, seqlen_k, is_prompt, full_len, is_prompt ? seq->cache.swa_valid_start : 0);
717                attn_flops += attn;
718                total_flops += total;
719            }
720
721            if (scheduler_outputs.min_seqlen_k == 0 || scheduler_outputs.min_seqlen_k > seqlen_k) {
722                scheduler_outputs.min_seqlen_k = seqlen_k;
723            }
724
725            seq_metadata_list[i] = std::move(p_metadata);
726            seq->schedule.start_inflight_step(max_seqlen_q);
727
728            i += 1;
729        }
730        scheduler_outputs.attn_tflops = attn_flops / 1e12;
731        scheduler_outputs.total_tflops = total_flops / 1e12;
732        last_step_num_batched_tokens = scheduler_outputs.num_batched_tokens;
733        last_step_cached_tokens = scheduler_outputs.num_cached_tokens_gpu;
734