[Kernel] RFC #29630 finale: retire sglang.jit_kernel into sglang.kernels (#32072)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py

@@ -17,7 +17,9 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python")
 import torch
 import triton
 
-from sglang.jit_kernel.cutedsl_kda import cutedsl_fused_sigmoid_gating_kda_update
+from sglang.kernels.ops.attention.cutedsl_kda import (
+    cutedsl_fused_sigmoid_gating_kda_update,
+)
 from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )

benchmark/bench_linear_attention/bench_gdn_qkv_split.py

@@ -4,7 +4,9 @@ import argparse
 
 import torch
 
-from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkv_split_gdn_prefill
+from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
+    fused_qkv_split_gdn_prefill,
+)
 
 DTYPES = {
     "bf16": torch.bfloat16,

benchmark/hicache/bench_hicache_write_back.py

@@ -12,7 +12,7 @@ from sgl_kernel.kvcacheio import (
     transfer_kv_all_layer_mla_lf_pf,
 )
 
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     can_use_hicache_jit_kernel,
     transfer_hicache_all_layer_mla_staged_lf_pf,
     transfer_hicache_all_layer_staged_lf_pf,
@@ -327,8 +327,8 @@ Environment:
 
 Comparison target:
 
-- MHA: `sgl_kernel.transfer_kv_all_layer_lf_pf` vs `sglang.jit_kernel.hicache.transfer_hicache_all_layer_staged_lf_pf`
-- MLA: `sgl_kernel.transfer_kv_all_layer_mla_lf_pf` vs `sglang.jit_kernel.hicache.transfer_hicache_all_layer_mla_staged_lf_pf`
+- MHA: `sgl_kernel.transfer_kv_all_layer_lf_pf` vs `sglang.kernels.ops.kvcache.hicache.transfer_hicache_all_layer_staged_lf_pf`
+- MLA: `sgl_kernel.transfer_kv_all_layer_mla_lf_pf` vs `sglang.kernels.ops.kvcache.hicache.transfer_hicache_all_layer_mla_staged_lf_pf`
 
 Metric:
 

benchmark/hicache/perf.py

@@ -15,7 +15,7 @@ def jit_hicache_impl(
     item_bytes: int,
     block_quota: int,
 ) -> None:
-    from sglang.jit_kernel.hicache import transfer_hicache_one_layer
+    from sglang.kernels.ops.kvcache.hicache import transfer_hicache_one_layer
 
     _ = item_bytes
 

benchmark/kernels/deepseek/benchmark_cute_dsl_fp8_paged_mqa_logits.py

@@ -8,8 +8,8 @@ import sys
 import numpy as np
 import torch
 
-import sglang.jit_kernel.dsa.cutedsl_paged_mqa_logits  # noqa: F401
-from sglang.jit_kernel.dsa import pick_dsl_expand
+import sglang.kernels.ops.attention.dsa.cutedsl_paged_mqa_logits  # noqa: F401
+from sglang.kernels.ops.attention.dsa import pick_dsl_expand
 from sglang.srt.layers.attention.dsa.utils import (
     fp8_mqa_logits_ceil_to_ue8m0,
     fp8_mqa_logits_make_fused_kv,