deepinfer-prefill
deepinfer/csrc/engine/scheduler.hpp · a679bb8 · L450–549 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
450
451 if (!prefix_caching_hit) {
452 /// Here it is assumed that the cache block size is larger
453 /// than the nextn num.
454 TORCH_CHECK(seq->cache.num_tokens == 0 || seq->cache.num_tokens >= this->lookahead_tokens,
455 SPRINT_VARS("Expected seq->cache.num_tokens == 0 || "
456 "seq->cache.num_tokens >= this->lookahead_tokens",
457 seq->cache.num_tokens, this->lookahead_tokens));
458 }
459
460 int seq_workload = num_prompt_tokens; /// 初始的 workload，用来判断是否 chunked
461 int new_workload = seq_workload; /// 施加各限制条件后的 workload
462 int seqlen_q = seq_workload;
463
464 TORCH_CHECK(seq_workload >= this->lookahead_tokens,
465 SPRINT_VARS("Expected seq_workload >= this->lookahead_tokens", seq_workload, this->lookahead_tokens));
466
467 // 满足下面的条件才会对 prompt 做截断 (截断到基础长度)：
468 // 加上当前这个 prompt 后会超过基础长度
469 /// @note 算剩余token时要考虑到cp的pad
470 if (seq_workload > (max_num_batched_tokens - num_batched_tokens) / cp_align * cp_align) {
471 new_workload = std::min(seq_workload, (max_num_batched_tokens - num_batched_tokens) / cp_align * cp_align);
472 seqlen_q = new_workload;
473 TORCH_CHECK(
474 seqlen_q < seq_workload, "_schedule: is_chunked must be true:, seqlen_q:", seqlen_q, " seq_workload: ", seq_workload);
475 }
476
477 auto start_offset = seq->cache.num_tokens;
478 while (!seq->schedule.chunk_checkpoints.empty() && seq->schedule.chunk_checkpoints.front() <= start_offset) {
479 seq->schedule.chunk_checkpoints.pop_front();
480 }
481 int checkpoint = seq->schedule.chunk_checkpoints.empty() ? 0 : seq->schedule.chunk_checkpoints.front();
482
483 if (start_offset + seqlen_q >= checkpoint && checkpoint > start_offset) {
484 seqlen_q = checkpoint - start_offset;
485 seq->schedule.chunk_checkpoints.pop_front();
486 }
487 // 调整多模态切的位置
488 int new_seqlen_q = _adjust_chunk_size(seqlen_q, seq->cache.num_tokens, seq->get_token_ids(), scheduled.empty());
489 if (num_batched_tokens + new_seqlen_q > max_num_batched_tokens) {
490 seqlen_q = _adjust_chunk_size(seqlen_q, seq->cache.num_tokens, seq->get_token_ids(), false);
491 } else {
492 seqlen_q = new_seqlen_q;
493 }
494 if (seqlen_q == 0) {
495 TORCH_CHECK(!scheduled.empty(), std::format("有请求正在排队，但是调度不上来，不应该 (post_adjust)"
496 ", len={}, cached={}, seq_workload={}, new_workload={}"
497 ", num_batched_tokens={}, max_num_batched_tokens={}"
498 ", cp_align={}, branch_id={}",
499 seq->get_len(), seq->cache.num_tokens, seq_workload, new_workload, num_batched_tokens,
500 max_num_batched_tokens, cp_align, seq->schedule.branch_id));
501 break;
502 }
503
504 bool is_chunked = seqlen_q < seq_workload;
505 if (is_chunked && !chunked_prefill) {
506 break;
507 }
508 // 其它约束都是硬约束，这里只验证 seqlen_q 满足最低要求，不满足则不调度
509 if (seqlen_q < std::max(lookahead_tokens, 1)) {
510 break;
511 }
512
513 // —— 确定调度该 seq，以下集中提交所有带副作用的操作 ——
514
515 num_batched_tokens += (seqlen_q + cp_align - 1) / cp_align * cp_align;
516
517 TORCH_CHECK(seq->schedule.status == SequenceStatus::WAITING, "status ", int(seq->schedule.status));
518 if (seq->cache.block_table.empty()) {
519 acquire.allocate_seq(*seq);
520 }
521
522 if (!is_chunked) {
523 num_curr_seqs += 1;
524 seq->schedule.status = SequenceStatus::RUNNING;
525 allocate_mtp_info_slot(*seq);
526 }
527
528 scheduled.push_back(seq);
529 TORCH_CHECK(seqlen_q + seq->cache.num_tokens <= seq->get_len(), SPRINT_VARS("Expected seqlen_q + seq->cache.num_tokens "
530 "<= seq->get_len()",
531 seqlen_q, seq->cache.num_tokens, seq->get_len()));
532 seqlens_q.push_back(seqlen_q);
533 if (is_chunked) {
534 auto current_prefill_len = seqlen_q + seq->cache.num_tokens;
535 seq->schedule.chunked_rest_len = seq->get_len() - current_prefill_len;
536
537 tokens_chunked_prefill_remaining += seq->get_len();
538 } else {
539 seq->schedule.chunked_rest_len = 0;
540 }
541 waiting.pop_front();
542
543 if (is_chunked && engine_type == EngineType::NORMAL) {
544 break;
545 }
546 }
547
548 if (engine_type != EngineType::NORMAL || !scheduled.empty() || !ignored_seqs.empty()) {
549 return _make_outputs(std::move(scheduled), true, num_batched_tokens, {}, std::move(ignored_seqs), std::move(seqlens_q), 0);