llm-inference-batching-scheduler

LLM 推理批处理调度器

此技能提供解决 LLM 推理批处理和调度优化问题的指导：需要将请求分组为批次，同时尽量降低成本、填充浪费和延迟。

问题理解

开始实现前，请彻底分析问题结构：

约束分析

1. 识别所有硬性约束——提取以下项目的具体限制：

   • 允许的最大唯一形状数量（例如，所有分桶合计 ≤ 8 种形状）
   • 延迟阈值（P95、P99）
   • 成本预算阈值
   • 填充率限制

2. 尽早计算硬性下界——开始编码前，先计算：

   • 对齐要求带来的最小可能填充量
   • 满足覆盖要求所需的最少批次数
   • 在约束条件下可达到的最高效率

3. 分解成本函数——了解各个组成部分：

   • 每批次开销（每批次的固定成本）
   • 形状编译成本（通常与序列长度呈二次关系）
   • 预填充／解码成本（每个请求的可变成本）
   • 按以下形式记录：Cost ≈ overhead × num_batches + shape_compile_cost + prefill_cost + decode_cost

数据分析

1. 分析请求分布——检查：

   • 提示词长度（prompt_len）的分布
   • 生成长度（gen_len）的分布
   • 找出可能对指标产生不成比例影响的异常值

2. 验证覆盖要求——确保：

   • 每个分桶中的最大 prompt_len 均由所选形状覆盖
   • 能处理 gen_len 极端值的边界情况

实现方法

构建可复用的评估基础设施

在调整参数之前，先创建系统化的评估工具：

1. Write a function that takes parameters (shape_list, gen_bucket_sizes) and returns all metrics
2. Include automatic constraint verification with assertions
3. Enable rapid parameter comparison without manual re-runs

参数搜索策略

避免随机试错，改用以下方法：

1. 小参数空间使用网格搜索——当参数范围有限时（例如，gen_bucket_size 在 [15, 50] 之间），系统地评估各种组合

2. 单参数使用二分搜索——固定其他参数、优化单个参数时，使用二分搜索找到最优值

3. 记录优化空间——跟踪各参数组合对应的指标值，以了解其中的权衡关系

形状选择指南

选择用于序列长度分桶的形状时：

1. 分析长度分布——选择能尽量减少最常见长度填充量的形状
2. 考虑 2 的幂或几何级数——这通常能在覆盖范围和形状数量之间取得平衡
3. 同时考虑所有分桶——如果各分桶共享形状，应进行全局优化，而非分别优化

生成长度分桶

gen_len 分桶参数会显著影响填充率：

1. 更小的分桶 = 填充率更低，但批次数更多（成本更高）
2. 更大的分桶 = 批次数更少，但 gen_len 的差异会带来更多填充
3. 通过计算填充预算并反向推导，找到最佳平衡点

验证策略

尽早检查约束

生成输出后立即验证：

1. All request IDs appear exactly once
2. Number of unique shapes ≤ limit
3. Each request's prompt_len ≤ assigned shape
4. No missing shapes that would cause alignment failures

指标验证

确认解决方案完整之前：

1. 运行官方评估脚本（如果提供）
2. 将所有指标与所有阈值进行比较
3. 分别检查每个分桶——一个分桶通过并不代表其他分桶也通过

常见验证失败

注意以下问题：

• 缺少形状导致覆盖不全（例如，需要形状 2048 时却未包含）
• 单请求批次浪费每批次开销
• 分别优化各分桶时违反形状约束

常见陷阱

过早优化

• 错误做法：尚未理解数学约束就开始实现
• 改进方法：提前花时间计算精确预算（例如，“bucket_1 最多可容忍 25,735 个填充 token”）

成本分析不足

• 错误：不了解哪项成本占主导
• 改进方法：在优化前计算并记录完整的成本明细

独立优化分桶

• 错误：约束同时涉及两个分桶，却分别优化每个分桶
• 改进方法：考虑联合优化，尤其是存在共享形状约束时

手动循环调参

• 错误：反复手动更改参数并重新运行
• 改进方法：编写参数扫描脚本，自动测试多种组合

太晚发现违反约束

• 错误：经过大量迭代后才检查约束
• 改进方法：在生成输出后立即添加断言检查

忽略基线实现

• 错误：未分析提供的基线代码以了解其低效之处
• 改进方法：研究 baseline_packer.py（或等效文件），找出具体导致其表现欠佳的原因

优化权衡

在优化前明确记录以下权衡：

调节项	降低	提高
更多形状	填充比例	形状编译成本
更小的 gen 分桶	填充比例	批次数量（成本）
更大的批量大小	每批开销	潜在填充量
更紧的形状间隔	填充量	所需形状数量

推荐工作流

1. 分析 - 解析问题，识别约束并计算硬性界限
2. 检测 - 构建评估工具，自动检查约束
3. 基线 - 运行基线方案，了解其不足
4. 拆解 - 将成本拆分为各个组成部分，找出主导项
5. 搜索 - 使用系统化的参数搜索（网格搜索／二分搜索），而非随机探索
6. 验证 - 检查所有分桶的约束和指标
7. 迭代 - 如果未通过，找出距离阈值最远的指标并重点优化