# vllm/v1/core/sched/scheduler.py

class Scheduler:
    def __init__(self, ...):
        # 等待队列：存放新到达和被抢占的请求
        self.waiting = create_request_queue(self.policy)

        # 运行队列：存放正在处理的请求
        self.running: list[Request] = []

        # 请求字典：通过 request_id 快速查找
        self.requests: dict[str, Request] = {}

class Scheduler:
    def __init__(self, ...):
        # 最大并发请求数
        self.max_num_running_reqs = self.scheduler_config.max_num_seqs

        # 每步最大 token 数
        self.max_num_scheduled_tokens = self.scheduler_config.max_num_batched_tokens

        # 最大序列长度
        self.max_model_len = vllm_config.model_config.max_model_len

# vllm/v1/core/sched/output.py

@dataclass
class SchedulerOutput:
    """调度器的输出，传递给 ModelExecutor"""

    # 新调度的请求
    scheduled_new_reqs: list[NewRequestData]

    # 继续运行的请求
    scheduled_cached_reqs: CachedRequestData

    # 每个请求调度的 token 数
    num_scheduled_tokens: dict[str, int]

    # 总调度 token 数
    total_num_scheduled_tokens: int

    # 抢占的请求 ID
    preempted_req_ids: set[str]

    # 完成的请求 ID
    finished_req_ids: set[str]

def schedule(self) -> SchedulerOutput:
    # ... 初始化 ...

    # 第一步：处理 RUNNING 请求
    req_index = 0
    while req_index < len(self.running) and token_budget > 0:
        request = self.running[req_index]

        # 计算需要调度的 token 数
        num_new_tokens = (
            request.num_tokens_with_spec      # 总 token 数
            + request.num_output_placeholders  # 输出占位符
            - request.num_computed_tokens      # 已计算的 token 数
        )

        # 应用长 prefill 分块限制
        if 0 < threshold < num_new_tokens:
            num_new_tokens = threshold

        # 不超过 token 预算
        num_new_tokens = min(num_new_tokens, token_budget)

        # 不超过模型最大长度
        num_new_tokens = min(
            num_new_tokens,
            self.max_model_len - 1 - request.num_computed_tokens
        )

        if num_new_tokens == 0:
            req_index += 1
            continue

        # 尝试分配 KV Cache
        while True:
            new_blocks = self.kv_cache_manager.allocate_slots(
                request,
                num_new_tokens,
                num_lookahead_tokens=self.num_lookahead_tokens,
            )

            if new_blocks is not None:
                # 分配成功
                break

            # 分配失败，执行抢占
            preempted_req = self._select_preempt_target()
            self._preempt_request(preempted_req, timestamp)

            if preempted_req == request:
                # 自己被抢占了，无法继续
                break

        if new_blocks is None:
            break

        # 调度成功
        scheduled_running_reqs.append(request)
        token_budget -= num_new_tokens
        req_index += 1

    # 第二步：处理 WAITING 请求（只有没有抢占时才处理）
    if not preempted_reqs:
        while self.waiting and token_budget > 0:
            # 检查并发限制
            if len(self.running) == self.max_num_running_reqs:
                break

            request = self.waiting.peek_request()

            # 检查各种等待状态
            if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
                # 等待远程 KV Cache 加载
                ...

            if request.status == RequestStatus.WAITING_FOR_FSM:
                # 等待语法编译
                ...

            # 查找前缀缓存
            new_computed_blocks, num_cached = (
                self.kv_cache_manager.get_computed_blocks(request)
            )

            # 计算需要调度的 token 数
            num_new_tokens = request.num_tokens - num_cached
            num_new_tokens = min(num_new_tokens, token_budget)

            # 分配 KV Cache
            new_blocks = self.kv_cache_manager.allocate_slots(
                request,
                num_new_tokens,
                num_new_computed_tokens=num_cached,
                new_computed_blocks=new_computed_blocks,
            )

            if new_blocks is None:
                # 无法分配，停止处理 waiting 队列
                break

            # 调度成功，移入 running 队列
            request = self.waiting.pop_request()
            self.running.append(request)
            request.status = RequestStatus.RUNNING
            request.num_computed_tokens = num_cached
            token_budget -= num_new_tokens

def _select_preempt_target(self) -> Request:
    """选择要抢占的请求"""
    if self.policy == SchedulingPolicy.PRIORITY:
        # 优先级调度：抢占优先级最低、到达最晚的请求
        return max(
            self.running,
            key=lambda r: (r.priority, r.arrival_time),
        )
    else:
        # FCFS：抢占最后加入的请求
        return self.running[-1]

def _preempt_request(self, request: Request, timestamp: float) -> None:
    """抢占请求并放回 waiting 队列"""
    assert request.status == RequestStatus.RUNNING

    # 释放 KV Cache
    self.kv_cache_manager.free(request)
    self.encoder_cache_manager.free(request)

    # 更新请求状态
    request.status = RequestStatus.PREEMPTED
    request.num_computed_tokens = 0  # 重置计算进度
    request.num_preemptions += 1

    # 放回 waiting 队列头部（优先恢复）
    self.waiting.prepend_request(request)

# 默认策略
class SchedulingPolicy(Enum):
    FCFS = "fcfs"
    PRIORITY = "priority"

# 基于优先级的调度
def _select_preempt_target(self):
    if self.policy == SchedulingPolicy.PRIORITY:
        # 选择优先级最低的请求
        return max(
            self.running,
            key=lambda r: (r.priority, r.arrival_time),
        )

# 长 prefill 分块阈值
threshold = self.scheduler_config.long_prefill_token_threshold

if 0 < threshold < num_new_tokens:
    num_new_tokens = threshold  # 限制每次处理的 token 数

def update_from_output(
    self,
    model_runner_output: ModelRunnerOutput,
    ...
) -> EngineCoreOutputs:
    """根据模型输出更新请求状态"""

    for req_id, sampler_output in model_output.items():
        request = self.requests[req_id]

        # 追加输出 token
        request.append_output_token_ids(sampler_output.sampled_token_ids)

        # 检查停止条件
        stopped = check_stop(request, self.max_model_len)

        if stopped:
            # 请求完成
            self._free_request(request)
            self.finished_req_ids.add(req_id)
            outputs.append(...)

# 增加最大并发数
llm = LLM(model="...", max_num_seqs=512)

# 增加每步 token 数
llm = LLM(model="...", max_num_batched_tokens=4096)

# 启用分块预填充
llm = LLM(
    model="...",
    enable_chunked_prefill=True,
    long_prefill_token_threshold=256,
)

# 使用优先级调度
llm = LLM(model="...", policy="priority")

# 发送请求时设置优先级
llm.generate(prompt, priority=0)  # 高优先级
llm.generate(prompt, priority=10)  # 低优先级