# PR #21483 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[misc] multiprocess compilation to speed up test
- 合并时间：2026-03-31 08:56
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/21483

---

# 执行摘要
该 PR 通过引入多进程并行编译 JIT 内核的机制，将自定义 AllReduce 测试时间从 300 秒优化至 150 秒，显著加速 CI 流水线。核心变更包括重构测试文件以支持预编译、调整内核标识符和测试参数，但 review 中提到的多进程启动方法鲁棒性问题未完全解决。

# 功能与动机
**动机**：根据 PR body 描述，主要目标是“Speed up JIT custom all reduce test. From 300s -> 150s.”，即减少 JIT 自定义 AllReduce 测试的执行时间，提升开发效率和 CI 性能。

# 实现拆解
实现涉及两个关键文件：
1. **`python/sglang/jit_kernel/all_reduce.py`**：修改 JIT 内核编译标识符，从 `custom_all_reduce` 重命名为 `custom_all_reduce_pull` 和 `custom_all_reduce_push`，以明确区分推拉操作。
2. **`python/sglang/jit_kernel/tests/test_custom_all_reduce.py`**：
 - 新增 `_precompile_kernels()` 函数，使用 `multiprocessing` 并行编译所有数据类型（如 float16, bfloat16）和世界大小（2-8）的内核组合。
 - 代码示例：
     ```python
     def _precompile_kernels() -> None:
         process_map: Dict[Tuple[torch.dtype, int], mp.Process] = {}
         COMPILE_SPACE = itertools.product(TEST_DTYPES, [2, 3, 4, 5, 6, 7, 8])
         mp.set_start_method("spawn")
         for config in COMPILE_SPACE:
             process_map[config] = mp.Process(target=_compile_one, args=config)
         for process in process_map.values():
             process.start()
         for (dtype, world_size), process in process_map.items():
             process.join()
             if process.exitcode != 0:
                 raise RuntimeError(f"Custom All Reduce {world_size=} {dtype=} failed")
     ```
 - 调整测试参数：`TEST_LAYERS` 从 2 增加到 4 以扩展测试覆盖；CI 时间估计从 500 秒降至 300 秒；优化 `pytest.mark.parametrize` 包含 nproc=1 的特殊处理。

# 评论区精华
review 中仅有一条高优先级评论：
> **gemini-code-assist[bot]**："Calling `mp.set_start_method("spawn")` can raise a `RuntimeError` if the start method has already been set... To make this more robust, you should handle the case where the start method is already been set."

建议添加 try-except 处理，但提交历史未显示采纳，可能被视为低风险或后续处理。

# 风险与影响
- **技术风险**：多进程编译可能引入竞争条件或资源冲突；内核标识符变更可能影响依赖代码；未处理 RuntimeError 可能导致测试在不稳定环境下失败。
- **影响范围**：主要优化内部测试流程，无直接用户影响；CI 测试时间减少 50%，提升开发迭代速度；可能增加测试环境的内存和 CPU 使用。

# 关联脉络
- 与 PR #21834（JIT rmsnorm 更新）同属 JIT 内核优化领域，可参考其性能提升模式。
- 与 PR #21783（TRT-LLM 稀疏 MLA 内核支持）在加速测试和内核调整方面有相似目标。
- 近期历史 PR 显示仓库持续关注测试优化和 CI 效率（如 #21873 添加网络超时、#21830 修复 CI 稳定性），本 PR 是这一趋势的延续。