「赤兔」Chitu 框架深度解读（十）：任务调度与并发控制策略

「赤兔」Chitu 框架深度解读（十）：任务调度与并发控制策略

高效的请求调度和并发控制是推理服务框架的核心能力，直接关系到系统的吞吐量、延迟和资源利用率。「赤兔」Chitu 框架设计了灵活的调度器 (scheduler.py) 和任务管理机制 (task.py) 来应对不同场景的需求，包括 KV Cache 管理、优先级处理、并发控制以及分布式环境下的协调。

核心组件：Task 与 TaskPool

• Task (task.py): 代表一个独立的推理请求。封装了请求信息 (req)、当前状态（task_type, waiting, consumed_req_tokens）、调度信息（priority, arrv_ts, sched_ddl, sched_group_id）、运行中间结果（next_token, handle）以及 KV Cache 占用信息。
• TaskPool (task.py): 全局的任务池，使用字典 (pool) 存储所有活跃的 Task，并维护一个 ID 列表 (id_list)。提供添加、移除、查询任务等基本操作。

Scheduler 基类与核心逻辑

Scheduler (scheduler.py) 是所有调度器的基类。其核心 schedule() 方法执行以下步骤：

1. 检查资源: 检查 TaskPool 是否为空，是否有空闲的调度组 (free_sgroups)。
2. 收集就绪任务: 从 TaskPool 中筛选出未处于等待状态 (waiting=False) 且需要模型运行 (no_model_run()=False) 的任务。
3. 严格模式过滤 (可选): 如果配置了 prefill_only 或 decode_only，则只保留对应类型的任务。
4. 排序: 根据 scorer() 方法计算每个任务的优先级得分，进行降序排序。scorer() 方法会根据配置组合多种评分策略（如 FCFS, Prefill优先, 请求优先级等）。
5. 确定任务类型: 以最高优先级任务的类型（Prefill 或 Decode）作为本轮调度的目标类型。
6. 调度 Prefill: 如果目标是 Prefill：
   • 调用 _schedule_prefill_tasks() 进行筛选。
   • Chunk Prefill: 如果启用了 prefill_chunk_size，根据预算切分 Prefill 任务，并更新 Task 的 prefill_chunk_size。
   • KV Cache 容量检查: 模拟分配 KV Cache Block，确保有足够空间容纳选中的 Prefill 任务。如果空间不足，则减少任务数量，直到满足条件。引入了 kvcache_block_threshold 进行动态容量控制。
   • 如果最终没有 Prefill 任务可选（例如 Cache 不足），则尝试调度 Decode 任务。
7. 调度 Decode: 如果目标是 Decode：
   • 调用 _schedule_decode_tasks() 进行筛选。
   • KV Cache 容量检查与驱逐: 检查是否有足够 Block 扩展 KV Cache。如果不足，则驱逐最低优先级的 Decode 任务（调用 evict_decode_task()），释放其 KV Cache，直到空间足够。被驱逐的任务状态会回滚到 Prefill 阶段。
8. 分配调度组: 从 free_sgroups 中取出一个 sgroup_id 分配给选中的任务，并更新 used_sgroups 和 sgroup_waiting_tasks。
9. 返回任务列表: 返回最终选定的任务 ID 列表。

特殊调度器实现

1. SkewScheduler

用于 cache_type="skew" 场景（可能是某种特定的 KV Cache 分片或调度策略）。

• 调度组: sgroup_id 对应 Cache 的 Slot ID。每个 Slot 有容量限制 (slot_handle.get_slot_size)。
• 调度逻辑: 优先选择最早释放的空闲 Slot (free_sgroups.popleft())。如果该 Slot 未满，则从外部任务池中选择最高优先级的、与 Slot 中现有任务类型一致的任务进行填充，直到 Slot 满或没有合适任务。然后调度该 Slot 中的所有任务（或部分任务，如果应用了 Chunk Prefill）。
• 任务移除: 当任务完成时，直接从对应的 sgroup_list 中移除，而不是简单标记。

2. DPFifoScheduler

用于 Expert Data Parallel (EDP) 场景，需要跨 DP Rank 协调调度。

• 目标: 为每个 DP Rank 分配不超过 max_num_tasks_per_dp 的任务列表。
• 调度逻辑:
  • 优先 Prefill: 先收集所有就绪的 Prefill 任务，按到达时间排序。
  • Cache Owner: 为首次出现的 Prefill 任务轮流分配一个 cache_owner (DP Rank ID)，后续调度沿用此 Owner，避免 KV Cache 跨 Rank 迁移。
  • 任务分配: 将 Prefill 任务根据其 cache_owner 分配到对应的 DP Rank 的任务列表。
  • Chunk Prefill (DP 感知): 如果启用，总的 prefill_chunk_size 会被分配到各个 DP Rank，然后每个 Rank 根据自己的预算对分配到的 Prefill 任务进行切分。
  • 数量截断: 确保每个 Rank 的任务数不超过 max_num_tasks_per_dp。
  • Prefill Fallback: 如果没有可选的 Prefill 任务（或 Chunk 切分后为空），则尝试调度 Decode 任务。
  • 调度 Decode: 收集就绪的 Decode 任务，根据其 cache_owner 分配到对应 Rank 的列表，并进行数量截断。
• 任务收集: 使用 DPTaskCollector.prepare_dp_tasks() 在所有 DP Rank 间同步最终的任务列表。

并发控制与 KV Cache 管理

• 调度组 (num_scheduler_groups): 限制了同时有多少批次的任务可以被调度和执行。在 PP 场景下，通常等于 PP size。在 Skew 场景下，等于 Slot 数量。
• Prefill/Decode 批大小限制: prefill_num_tasks 和 decode_num_tasks 直接控制了每个调度组内 Prefill 和 Decode 任务的最大数量。
• KV Cache 容量:
  • Prefill 调度会进行严格的容量检查，避免 OOM。
  • Decode 调度在容量不足时会主动驱逐低优先级任务，保证高优先级任务的运行，这是一种动态的拥塞控制机制。
  • 引入 kvcache_block_threshold，在发生驱逐时会降低此阈值，而在任务正常完成时会恢复，动态调整系统的接纳能力。

总结

「赤兔」的调度器设计兼顾了灵活性和鲁棒性。通过组合不同的评分策略，可以实现多样的优先级调度。针对 Prefill 和 Decode 阶段的不同特点（Prefill 消耗大量新 Block，Decode 逐步扩展 Block），设计了不同的容量检查和处理逻辑。特别是在 Decode 阶段引入的驱逐机制，保证了即使在 KV Cache 接近饱和时系统也能继续服务高优先级请求。针对 Skew Cache 和 DP 并行等特殊场景，也提供了定制化的调度器实现，确保任务能被正确且高效地分配到相应的计算资源或 Cache 分片上。# 「赤兔」Chitu 框架深度解读（十）：任务调度与并发控制策略

高效的请求调度和并发控制是推理服务框架的核心能力，直接关系到系统的吞吐量、延迟和资源利用率。「赤兔」Chitu 框架设计了灵活的调度器 (scheduler.py) 和任务管理机制 (task.py) 来应对不同场景的需求，包括 KV Cache 管理、优先级处理、并发控制以及分布式环境下的协调。

核心组件：Task 与 TaskPool

• Task (task.py): 代表一个独立的推理请求。封装了请求信息 (req)、当前状态（task_type, waiting, consumed_req_tokens）、调度信息（priority, arrv_ts, sched_ddl, sched_group_id）、运行中间结果（next_token, handle）以及 KV Cache 占用信息。
• TaskPool (task.py): 全局的任务池，使用字典 (pool) 存储所有活跃的 Task，并维护一个 ID 列表 (id_list)。提供添加、移除、查询任务等基本操作。

Scheduler 基类与核心逻辑

Scheduler (scheduler.py) 是所有调度器的基类。其核心 schedule() 方法执行以下步骤：

1. 检查资源: 检查 TaskPool 是否为空，是否有空闲的调度组 (free_sgroups)。
2. 收集就绪任务: 从 TaskPool 中筛选出未处于等待状态 (waiting=False) 且需要模型运行 (no_model_run()=False) 的任务。
3. 严格模式过滤 (可选): 如果配置了 prefill_only 或 decode_only，则只保留对应类型的任务。
4. 排序: 根据 scorer() 方法计算每个任务的优先级得分，进行降序排序。scorer() 方法会根据配置组合多种评分策略（如 FCFS, Prefill优先, 请求优先级等）。
5. 确定任务类型: 以最高优先级任务的类型（Prefill 或 Decode）作为本轮调度的目标类型。
6. 调度 Prefill: 如果目标是 Prefill：
   • 调用 _schedule_prefill_tasks() 进行筛选。
   • Chunk Prefill: 如果启用了 prefill_chunk_size，根据预算切分 Prefill 任务，并更新 Task 的 prefill_chunk_size。
   • KV Cache 容量检查: 模拟分配 KV Cache Block，确保有足够空间容纳选中的 Prefill 任务。如果空间不足，则减少任务数量，直到满足条件。引入了 kvcache_block_threshold 进行动态容量控制。
   • 如果最终没有 Prefill 任务可选（例如 Cache 不足），则尝试调度 Decode 任务。
7. 调度 Decode: 如果目标是 Decode：
   • 调用 _schedule_decode_tasks() 进行筛选。
   • KV Cache 容量检查与驱逐: 检查是否有足够 Block 扩展 KV Cache。如果不足，则驱逐最低优先级的 Decode 任务（调用 evict_decode_task()），释放其 KV Cache，直到空间足够。被驱逐的任务状态会回滚到 Prefill 阶段。
8. 分配调度组: 从 free_sgroups 中取出一个 sgroup_id 分配给选中的任务，并更新 used_sgroups 和 sgroup_waiting_tasks。
9. 返回任务列表: 返回最终选定的任务 ID 列表。

特殊调度器实现

1. SkewScheduler

用于 cache_type="skew" 场景（可能是某种特定的 KV Cache 分片或调度策略）。

• 调度组: sgroup_id 对应 Cache 的 Slot ID。每个 Slot 有容量限制 (slot_handle.get_slot_size)。
• 调度逻辑: 优先选择最早释放的空闲 Slot (free_sgroups.popleft())。如果该 Slot 未满，则从外部任务池中选择最高优先级的、与 Slot 中现有任务类型一致的任务进行填充，直到 Slot 满或没有合适任务。然后调度该 Slot 中的所有任务（或部分任务，如果应用了 Chunk Prefill）。
• 任务移除: 当任务完成时，直接从对应的 sgroup_list 中移除，而不是简单标记。

2. DPFifoScheduler

用于 Expert Data Parallel (EDP) 场景，需要跨 DP Rank 协调调度。

• 目标: 为每个 DP Rank 分配不超过 max_num_tasks_per_dp 的任务列表。
• 调度逻辑:
  • 优先 Prefill: 先收集所有就绪的 Prefill 任务，按到达时间排序。
  • Cache Owner: 为首次出现的 Prefill 任务轮流分配一个 cache_owner (DP Rank ID)，后续调度沿用此 Owner，避免 KV Cache 跨 Rank 迁移。
  • 任务分配: 将 Prefill 任务根据其 cache_owner 分配到对应的 DP Rank 的任务列表。
  • Chunk Prefill (DP 感知): 如果启用，总的 prefill_chunk_size 会被分配到各个 DP Rank，然后每个 Rank 根据自己的预算对分配到的 Prefill 任务进行切分。
  • 数量截断: 确保每个 Rank 的任务数不超过 max_num_tasks_per_dp。
  • Prefill Fallback: 如果没有可选的 Prefill 任务（或 Chunk 切分后为空），则尝试调度 Decode 任务。
  • 调度 Decode: 收集就绪的 Decode 任务，根据其 cache_owner 分配到对应 Rank 的列表，并进行数量截断。
• 任务收集: 使用 DPTaskCollector.prepare_dp_tasks() 在所有 DP Rank 间同步最终的任务列表。

并发控制与 KV Cache 管理

• 调度组 (num_scheduler_groups): 限制了同时有多少批次的任务可以被调度和执行。在 PP 场景下，通常等于 PP size。在 Skew 场景下，等于 Slot 数量。
• Prefill/Decode 批大小限制: prefill_num_tasks 和 decode_num_tasks 直接控制了每个调度组内 Prefill 和 Decode 任务的最大数量。
• KV Cache 容量:
  • Prefill 调度会进行严格的容量检查，避免 OOM。
  • Decode 调度在容量不足时会主动驱逐低优先级任务，保证高优先级任务的运行，这是一种动态的拥塞控制机制。
  • 引入 kvcache_block_threshold，在发生驱逐时会降低此阈值，而在任务正常完成时会恢复，动态调整系统的接纳能力。

总结

「赤兔」的调度器设计兼顾了灵活性和鲁棒性。通过组合不同的评分策略，可以实现多样的优先级调度。针对 Prefill 和 Decode 阶段的不同特点（Prefill 消耗大量新 Block，Decode 逐步扩展 Block），设计了不同的容量检查和处理逻辑。特别是在 Decode 阶段引入的驱逐机制，保证了即使在 KV Cache 接近饱和时系统也能继续服务高优先级请求。针对 Skew Cache 和 DP 并行等特殊场景，也提供了定制化的调度器实现，确保任务能被正确且高效地分配到相应的计算资源或 Cache 分片上。