执行摘要

• 一句话：路由concat_mla到JIT，移除未使用的downcast_fp8
• 推荐动作：建议合并。这是一个干净的代码清理，经过充分测试，无回归风险。

功能与动机

concat_mla已有JIT实现，CUDA运行时可直接使用，减少对重复AOT CUDA ops的依赖。downcast_fp8无运行时调用点，可安全删除。

实现拆解

1. 修改DeepSeek模型前向导入：在forward_mha.py中，将CUDA路径的concat_mla_k和merge_state_v2导入拆分，concat_mla_k改从sglang.jit_kernel.concat_mla导入；增加elif _is_musa分支保留从sgl_kernel导入。
2. 修改Sarvam MoE模型导入：在sarvam_moe.py中，类似地调整concat_mla_k的导入源，在CUDA路径下从JIT导入，MUSA路径保持不变。
3. 修改注意力工具函数导入：在utils.py中，将concat_mla_absorb_q的导入从sgl_kernel改为从JIT。
4. 删除无用代码：删除cast.py（包含downcast_fp8 JIT包装器）、bench_cast.py（基准测试文件）和elementwise/cast.cuh（CUDA头文件），这些代码已无使用。
5. 验证：通过compileall检查编译，运行test_concat_mla.py（94个测试通过），并在H200上验证。

关键文件：

• python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py（模块 注意力前向；类别 source；类型 dependency-wiring；符号 concat_mla_k）: DeepSeek系列模型注意力核心前向方法，导入变更直接影响运行时行为。
• python/sglang/srt/models/sarvam_moe.py（模块 MoE模型；类别 source；类型 dependency-wiring；符号 concat_mla_k）: Sarvam MoE模型也使用concat_mla_k，导入调整保持一致。
• python/sglang/srt/layers/attention/utils.py（模块 注意力工具；类别 source；类型 dependency-wiring；符号 concat_mla_absorb_q）: 注意力工具函数中的concat_mla_absorb_q导入也迁移到JIT，影响其他注意力计算路径。
• python/sglang/jit_kernel/cast.py（模块 JIT内核；类别 source；类型 deletion；符号 _jit_cast_module, downcast_fp8）: 删除未使用的downcast_fp8 JIT包装器及其辅助模块加载函数。
• python/sglang/jit_kernel/benchmark/bench_cast.py（模块 基准测试；类别 source；类型 deletion；符号 benchmark, _report_bandwidth, fmt, report_bandwidth）: 删除downcast_fp8的基准测试文件，因核心函数已删除。
• python/sglang/jit_kernel/csrc/elementwise/cast.cuh（模块 CUDA头文件；类别 other；类型 deletion）: 删除downcast_fp8的CUDA kernel头文件，核心实现已无用途。

关键符号：downcast_fp8, _jit_cast_module

关键源码片段

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

DeepSeek系列模型注意力核心前向方法，导入变更直接影响运行时行为。

# forward_mha.py 中 CUDA/MUSA 分支导入调整
if _is_cuda:
    from sgl_kernel import merge_state_v2 # 保留 AOT 的 merge_state_v2
    from sglang.jit_kernel.concat_mla import concat_mla_k # concat_mla_k 改从 JIT 导入
elif _is_musa:
    from sgl_kernel import concat_mla_k # MUSA 平台仍使用 sgl_kernel AOT 实现

python/sglang/srt/models/sarvam_moe.py

Sarvam MoE模型也使用concat_mla_k，导入调整保持一致。

# sarvam_moe.py 中 CUDA 分支导入调整
if _is_cuda:
    try:
        from sgl_kernel import bmm_fp8, merge_state_v2 # 保留 AOT 的其他符号
        from sglang.jit_kernel.concat_mla import concat_mla_k # concat_mla_k 从 JIT 导入
        ...

python/sglang/jit_kernel/cast.py

删除未使用的downcast_fp8 JIT包装器及其辅助模块加载函数。

# 已删除的 downcast_fp8 JIT 包装器（有专门的 JIT kernel 实现，无需此 Python 层包装）
# 该函数在运行时没有 import 或调用点，因此被整体移除。
@cache_once
def _jit_cast_module(dtype: torch.dtype) -> Module:
    args = make_cpp_args(dtype)
    return load_jit('cast', *args, cuda_files=['elementwise/cast.cuh'],
                    cuda_wrappers=[('downcast_fp8', f'downcast_fp8<{args}>')])
​
def downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc, mult=1, offset=0):
    # Fused downcast of KV cache tensors from bf16/fp16 to fp8 (E4M3).
    module = _jit_cast_module(k.dtype)
    module.downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset)

评论区精华

无实质性review讨论。作者在PR评论中报告了远程H200验证结果：Nsight Compute检查通过，编译和导入检查正常，JIT测试94个全部通过。

• 暂无高价值评论线程

风险与影响

• 风险：主要风险是JIT实现可能与AOT实现存在行为差异，但现有JIT测试覆盖充分（94个测试通过），且通过导入检查确认运行时路径正确。风险较低。删除的downcast_fp8代码已无引用，无回归风险。
• 影响：对用户：无功能变化，性能不受影响（JIT实现应与AOT一致）。对系统：减少依赖项，简化部署。对团队：降低维护成本，需注意未来新增CUDA调用时统一使用JIT路径。
• 风险标记：核心路径变更, 删除文件, 依赖项迁移

关联脉络

• PR #26000 [codex] Centralize Triton utility kernels: 同一文件utils.py中的导入重构，与本PR的JIT路由方向一致，属于同一代码清洁系列。