执行摘要

• 一句话：批量迁移 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 的典型模式

# 文件顶部导入块的变化（新旧对比虚线分隔）
​
# ====== 旧路径 (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

# 在 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 算子的导入替换

# 文件顶部导入变化
​
# 旧 : 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 基础设施的迁移框架。