# PR #32015 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[Kernel] Phase 4 batch-2: migrate JIT operator groups into kernels.ops (no shims) (RFC #29630)
- 合并时间：2026-07-22 17:49
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/32015

---

# 执行摘要

- 一句话：批量迁移 JIT 算子至 kernels.ops 命名空间
- 推荐动作：建议阅读该 PR 的组织方式和分类规则，了解 sglang 内核的命名空间规划。对于开发者，新算子应直接添加到 `sglang.kernels.ops` 下对应组中。

# 功能与动机

实现 sglang 内核命名空间的统一规划，消除路径碎片，提升代码可维护性。RFC #29630 定义了 Phase 4 的迁移计划，batch-2 为直接导入的 JIT 算子组。

# 实现拆解

1. **建立新目录结构**：在 `sglang/kernels/ops/` 下创建 `gemm`、`communication`、`layernorm`、`speculative`、`mamba`、`quantization`、`kvcache` 等子包，将原 `sglang/jit_kernel/` 中的对应算子文件迁移到新路径。
2. **更新所有调用点**：扫描仓库中直接 `from sglang.jit_kernel.xxx import yyy` 的代码，改为 `from sglang.kernels.ops.<group>.xxx import yyy`。涉及 `srt/mem_cache/pool_host/mha.py`、`srt/models/inkling_common/kernels/sconv.py`、`srt/layers/layernorm.py`、`srt/mem_cache/memory_pool_host.py`、`srt/models/minimax_m2.py` 等 95 个源文件。
3. **删除旧路径 shim**：移除 `sglang/jit_kernel/` 下的原算子文件，不保留任何向后兼容的 re-export，强制所有引用使用新路径。
4. **测试验证**：运行 68 个 CPU 内核测试用例全部通过，确保导入路径迁移正确。
5. **代码风格清理**：执行 lint 检查，保证无额外代码风格问题。

关键文件：
- `python/sglang/srt/mem_cache/pool_host/mha.py`（模块 缓存层；类别 source；类型 dependency-wiring）: 展示 hicache JIT 算子从 `sglang.jit_kernel.hicache` 迁移到 `sglang.kernels.ops.kvcache.hicache` 的典型模式
- `python/sglang/srt/models/inkling_common/kernels/sconv.py`（模块 Mamba 算子；类别 source；类型 data-contract）: 展示 mamba 算子从 `sglang.jit_kernel.inkling_sconv` 迁移到 `sglang.kernels.ops.mamba.inkling_sconv`
- `python/sglang/srt/layers/layernorm.py`（模块 LayerNorm；类别 source；类型 dependency-wiring）: 展示 layernorm 算子的导入替换
- `python/sglang/srt/models/minimax_m2.py`（模块 模型；类别 source；类型 data-contract）: 展示 all_reduce 算子的导入替换
- `python/sglang/srt/mem_cache/memory_pool_host.py`（模块 内存池；类别 source；类型 dependency-wiring）: 展示多个算子迁移，包括 hicache, hisparse, transfer_mamba

关键符号：未识别

## 关键源码片段

### `python/sglang/srt/mem_cache/pool_host/mha.py`

展示 hicache JIT 算子从 `sglang.jit_kernel.hicache` 迁移到 `sglang.kernels.ops.kvcache.hicache` 的典型模式

```python
# 文件顶部导入块的变化（新旧对比虚线分隔）

# ====== 旧路径 (Phase 4 batch-2 前 ) ======
# from sglang.jit_kernel.hicache import (
# can_use_hicache_jit_kernel,
# can_use_write_back_jit_kernel,
# )
# from sglang.jit_kernel.hicache import (
# transfer_hicache_all_layer as jit_transfer_hicache_all_layer,
# )
# ... 共 7 个导入语句分别引入不同函数

# ====== 新路径 (Phase 4 batch-2 后 ) ======
from sglang.kernels.ops.kvcache.hicache import (
    can_use_hicache_jit_kernel,
    can_use_write_back_jit_kernel,
)
from sglang.kernels.ops.kvcache.hicache import (
    transfer_hicache_all_layer as jit_transfer_hicache_all_layer,
)
# ... 其余导入类似，模块路径统一从 `sglang.jit_kernel.hicache` 变为 `sglang.kernels.ops.kvcache.hicache`

# 注意：所有别名（jit_transfer_*）保持不变，仅源路径迁移。

```

### `python/sglang/srt/models/inkling_common/kernels/sconv.py`

展示 mamba 算子从 `sglang.jit_kernel.inkling_sconv` 迁移到 `sglang.kernels.ops.mamba.inkling_sconv`

```python
# 在 causal_conv1d 函数中，CUDA 分支的导入变化

def causal_conv1d(...):
    # ... 前置检查
    if (
        is_cuda()
        and not is_decode
        and x.dtype == torch.bfloat16
        and D % 2 == 0
        and x.stride(1) == 1
    ):
        # 旧 : from sglang.jit_kernel.inkling_sconv import causal_conv1d as _cuda_causal_conv1d
        # 新 : from sglang.kernels.ops.mamba.inkling_sconv import (
        # causal_conv1d as _cuda_causal_conv1d,
        # )
        from sglang.kernels.ops.mamba.inkling_sconv import (
            causal_conv1d as _cuda_causal_conv1d,
        )
        return _cuda_causal_conv1d(
            x, weight, sconv_cache, cache_mask, safe_idx, cu, si,
            activation=activation, use_residual=use_residual, is_decode=is_decode,
        )
    # ... 剩余 Triton fallback 代码不变

```

### `python/sglang/srt/models/minimax_m2.py`

展示 all_reduce 算子的导入替换

```python
# 文件顶部导入变化

# 旧 : from sglang.jit_kernel.all_reduce import (
# fused_parallel_qknorm,
# get_fused_parallel_qknorm_max_occupancy,
# )
# 新 :
from sglang.kernels.ops.communication.all_reduce import (
    fused_parallel_qknorm,
    get_fused_parallel_qknorm_max_occupancy,
)

# 注意：同时调整了 import 顺序，`kernel_api_logging` 移至新导入之前

```

# 评论区精华

PR 无实质性人工讨论，仅包含自动机器人评论（Gemini Code Assist 已停止服务）和 CI 链接。合并前未引发设计争议。

- 暂无高价值评论线程

# 风险与影响

- 风险：风险较低，但 95 个文件的大规模导入路径变更可能遗漏个别调用点。测试覆盖了 CPU 内核，未覆盖 GPU kernel 执行路径；删除 shim 可能影响外部直接导入旧路径的脚本。需确认下游无直接依赖 `sglang.jit_kernel` 子模块。
- 影响：影响所有使用 `sglang.jit_kernel` 中 JIT 算子的内部模块。团队需遵循新命名规范编写新算子。外部用户若直接 `import sglang.jit_kernel.xxx` 将失效，但通常不建议直接导入内部路径。无功能变更，性能无影响。
- 风险标记：大规模导入路径变更 , 删除兼容 shim, 跨 95 个文件重构

# 关联脉络

- PR #31666 [Kernel] Phase 3+4: move JIT infra + operator groups into sglang.kernels (RFC #29630): 这是 Phase 4 batch-2 的前序 PR，完成了 JIT 基础设施的迁移框架。