执行摘要

本 PR 实现了 JIT 编译的激活内核，替换 CUDA 平台的 AOT 实现，旨在通过 PDL 和向量化优化提升 SiLU、GELU 等激活函数的性能，影响多个核心模块，但 review 中识别出 HIP 平台兼容性和输出形状注册问题尚未解决。

功能与动机

动机是优化推理速度，通过 JIT 编译在运行时生成高效内核。PR body 中提供了 H200 和 B200 GPU 的基准测试图，显示性能增益主要来自 PDL 和向量化，目的是减少延迟并统一内核实现。

实现拆解

实现按模块拆解如下：

• JIT 内核核心：新增 activation.cuh CUDA 文件，使用模板和 PDL 实现向量化激活核函数。
  

  template <typename T, ActivationKind kAct, bool kUsePDL>
  __global__ void act_and_mul_kernel(...) {
      // PDL 等待和触发，向量化加载/存储
  }
  

• Python 包装层：新增 activation.py，提供 run_activation 函数和类型化接口，通过 register_custom_op 集成到系统。
• 模块集成：修改了多个文件以替换导入，例如 srt/layers/activation.py 从 sgl_kernel 切换到 sglang.jit_kernel.activation。
• 测试与文档：新增单元测试验证正确性，基准测试对比 AOT/JIT/torch.compile 性能，并更新技能文档添加示例。

评论区精华

review 讨论聚焦于两个关键问题：

1. HIP 平台导入缺失：gemini-code-assist[bot] 指出：“The imports for the HIP platform are missing gelu_and_mul and gelu_tanh_and_mul... running this code on an AMD GPU will result in a NameError。” BBuf 回应询问是否应用建议，但未明确解决。
2. 输出形状注册错误：同一 review 指出：“out_shape="input" parameter ... is likely incorrect because this operation changes the tensor shape”，建议使用更准确的形状描述。

风险与影响

风险：

• HIP 平台导入缺失可能导致 AMD GPU 运行时崩溃，需添加缺失导入。
• 输出形状注册错误可能影响 torch.compile 或自动输出分配，引发形状不匹配错误。
• 核心路径变更引入回归风险，依赖测试覆盖确保数值正确性。
• PDL 优化在非支持架构上可能降级，需架构检测。

影响：

• 用户：CUDA 平台用户获得性能提升，但 AMD 用户可能受影响。
• 系统：统一了激活实现，减少代码冗余，但增加了 JIT 编译开销。
• 团队：需学习 JIT 内核开发和 PDL 使用，文档更新提供了指导。

关联脉络

与历史 PR 关联紧密：PR 22078 回滚了相同的 JIT 激活功能，原因是 CI 测试失败（见 Issue 评论）。本 PR 是重新引入，表明团队在解决稳定性问题后继续推进性能优化。这揭示了 sglang 仓库在 JIT 内核演进中的反复尝试，以及性能与稳定性之间的权衡。