deepinfer-commit
deepinfer/csrc/engine/scheduler.hpp · a679bb8 · L899–1003 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
899 for (ssize_t i = 0; i < num_output_tokens; i++) {
900 auto &seq = *scheduler_outputs.scheduled_seqs[i];
901 auto &seq_ptr = scheduler_outputs.scheduled_seqs[i];
902
903 seq.schedule.finish_inflight_step();
904
905 // seq is already aborted
906 if (seq.is_finished()) {
907 release_mtp_info_slot(seq);
908 pushed_tokens_before_finish.push_back(0);
909 if (scheduler_outputs.prompt_run) {
910 running.push_back(seq_ptr);
911 }
912 if (has_accepted_tokens)
913 flat_offset += scheduler_outputs.seqlens_q[i];
914 continue;
915 }
916
917 // ZOS 异步：本步尚在途时就被 _preempt 抢占重算的 seq（只在 decode 步出现——prompt 步的
918 // 在途 seq 不在 running 里，不会成为抢占目标）。_preempt 已复位其计数并放回 waiting 队头，
919 // 本步输出作废：跳过全部记账，避免重复入队与 cache.num_tokens 的错误推进。
920 if (!scheduler_outputs.prompt_run && seq.schedule.status == SequenceStatus::WAITING) {
921 pushed_tokens_before_finish.push_back(0);
922 if (has_accepted_tokens)
923 flat_offset += scheduler_outputs.seqlens_q[i];
924 continue;
925 }
926
927 seq.cache.advance_optimistically(scheduler_outputs.seqlens_q[i]);
928 if (so_status_bool)
929 seq.structured_output.status = so_status[i].cast<int>();
930
931 // 乐观推进后 cache.num_tokens 的不变式（input.tokens 只含已提交 token，不含 draft，契约见其声明处；
932 // 在途被抢占的 seq 已在上面跳过）：
933 // prompt 步：seqlen_q 只覆盖 prompt 本身，推进后 <= len，等于 len 说明本 seq 的 prefill 已完成；
934 // decode 步：seqlen_q = 1 + 参与验证的 draft 数，推进后其值落在 [len, len + nextn] 内，
935 // 多算的部分随后按实际接受的 token 数（token_cnt）回填修正。
936 TORCH_CHECK(seq.cache.num_tokens <= seq.input.tokens.get_len() + seq.get_nextn_len());
937 bool prefill_finish_or_is_decode = (seq.cache.num_tokens >= seq.input.tokens.get_len());
938 TORCH_CHECK(scheduler_outputs.prompt_run || prefill_finish_or_is_decode, "decode 步推进后必然 >= len，不应该：",
939 seq.cache.num_tokens, " < ", seq.input.tokens.get_len());
940 if (prefill_finish_or_is_decode) {
941 if (engine_type == EngineType::NORMAL) { /// prefill_engine won't change status
942 /// because it will be transferred to
943 /// complete_engine
944 TORCH_CHECK(seq_ptr->schedule.status == SequenceStatus::RUNNING,
945 SPRINT_VARS("`seq_ptr->schedule.status` should be RUNNING. "
946 "`seq_ptr->schedule.status` is updated in `_allocate`. "
947 "Check if the `allocate` is called at the end of "
948 "the prefill phase. Got",
949 int(seq_ptr->schedule.status)));
950 }
951 if (scheduler_outputs.prompt_run) {
952 running.push_back(seq_ptr);
953 }
954
955 // 提取本步验证通过的 draft tokens（flat 布局见 accepted_token_ids 参数注释）
956 accepted_token_ids_vec.clear();
957 if (has_accepted_tokens) {
958 int num_draft_i = scheduler_outputs.seqlens_q[i] - 1;
959 for (int j = 0; j < num_draft_i; j++) {
960 int token_id = accepted_token_ids_ptr[flat_offset + j];
961 if (token_id != -1) {
962 accepted_token_ids_vec.push_back(token_id);
963 } else
964 break;
965 }
966 }
967 seq.mtp.accepted_tokens = accepted_token_ids_vec;
968
969 token_ids = accepted_token_ids_vec;
970 seq.mtp.last_output_token = output_token_ids_data[i];
971 token_ids.push_back(seq.mtp.last_output_token);
972
973 int token_cnt = 0;
974 // 判死看账本 structured_output.status 而非当步上报：sampler 对批内无活跃 SO 的步返回 None，
975 // 只认上报会让账本已是 -1 的 seq（如 PE 判死后迁到 CE 的副本）无约束自由生成
976 if (seq.structured_output.status == -1) {
977 seq.schedule.status = SequenceStatus::FINISHED_SO_FAILED;
978 seq.append_token_id(0);
979 seq.append_compressed_token_id(0);
980 token_cnt++;
981 } else {
982 for (auto token_id : token_ids) {
983 seq.append_token_id(token_id);
984 seq.append_compressed_token_id(map_engram(token_id));
985 token_cnt++;
986 _check_stop(seq);
987 if (seq.is_finished()) {
988 break;
989 }
990 }
991 }
992 if (seq.is_finished()) {
993 release_mtp_info_slot(seq);
994 }
995 pushed_tokens_before_finish.push_back(token_cnt);
996
997 if (!scheduler_outputs.prompt_run) {
998 // decode：把 cache.num_tokens 从乐观推进值回填为实际接受数（未被接受的 draft 不占 KV）
999 seq.cache.commit_processed_tokens(token_cnt);
1000 }
1001 TORCH_CHECK(seq.cache.num_tokens <= seq.input.tokens.get_len(),
1002 "assert seq.cache.num_tokens <= seq.input.tokens.get_len(), seq.cache.num_tokens:", seq.cache.num_tokens,
1003 " seq.input.tokens.get_len():", seq.input.tokens.get_len());