[Kernel] Phase 4 batch-2: migrate JIT operator groups into kernels.ops (no shims) (RFC #29630) (#32015)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

test/manual/layers/moe/bench_mxfp4_sm90_kernels.py

@@ -36,7 +36,7 @@ from flashinfer.fused_moe import (
 from flashinfer.fused_moe.core import ActivationType
 
 # ---- SGLang Marlin ----
-from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
 from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import fused_marlin_moe
 from sglang.srt.layers.quantization.marlin_utils import (
     marlin_make_workspace,

test/registered/jit/benchmark/bench_awq_dequantize.py

@@ -4,8 +4,10 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
 from sglang.jit_kernel.benchmark.utils import run_benchmark
+from sglang.kernels.ops.quantization.awq_dequantize import (
+    awq_dequantize as jit_awq_dequantize,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/benchmark/bench_custom_all_reduce.py

@@ -12,8 +12,8 @@ import torch.distributed as dist
 import sglang.srt.distributed.parallel_state as ps
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, multigpu_bench_main
-from sglang.jit_kernel.mp import register_comm_cleanup
 from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl
+from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py

@@ -9,11 +9,11 @@ import torch.nn.functional as F
 import triton.testing
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import (
-    dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm,
-)
 from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
 from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import (
+    dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py

@@ -6,7 +6,7 @@ import torch
 import triton
 
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
-from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
+from sglang.kernels.ops.gemm.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
 from sglang.srt.utils import is_sm120_supported
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/benchmark/bench_fused_eh_norm.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import torch
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.fused_eh_norm import fused_eh_norm
+from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py

@@ -1,7 +1,7 @@
 import torch
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.fused_fp8_qkv_kv_cache import fused_fp8_qkv_kv_cache
+from sglang.kernels.ops.kvcache.fused_fp8_qkv_kv_cache import fused_fp8_qkv_kv_cache
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_hicache.py

@@ -26,7 +26,7 @@ import triton.testing
 from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
 
 from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     can_use_hicache_jit_kernel,
     transfer_hicache_all_layer,
     transfer_hicache_one_layer,

test/registered/jit/benchmark/bench_hisparse.py

@@ -6,7 +6,7 @@ import triton
 import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
-from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla
+from sglang.kernels.ops.kvcache.hisparse import load_cache_to_device_buffer_mla
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_mxfp8_moe.py

@@ -7,7 +7,7 @@ import torch
 import triton
 
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
-from sglang.jit_kernel.mxfp8 import (
+from sglang.kernels.ops.quantization.mxfp8 import (
     es_sm100_mxfp8_blockscaled_grouped_quant,
     es_sm100_mxfp8_blockscaled_moe_grouped_gemm,
 )

test/registered/jit/benchmark/bench_ngram_compute_decode.py

@@ -7,7 +7,7 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark_no_cudagraph,
 )
-from sglang.jit_kernel.ngram_embedding import (
+from sglang.kernels.ops.speculative.ngram_embedding import (
     compute_n_gram_ids,
     compute_n_gram_ids_decode,
 )

test/registered/jit/benchmark/bench_ngram_update_token_table.py

@@ -7,7 +7,7 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark_no_cudagraph,
 )
-from sglang.jit_kernel.ngram_embedding import (
+from sglang.kernels.ops.speculative.ngram_embedding import (
     update_token_table,
     update_token_table_decode,
 )

test/registered/jit/benchmark/bench_resolve_future_token_ids.py

@@ -9,7 +9,9 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
-from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda
+from sglang.kernels.ops.speculative.resolve_future_token_ids import (
+    resolve_future_token_ids_cuda,
+)
 from sglang.srt.utils import get_compiler_backend
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/benchmark/bench_rope.py

@@ -114,8 +114,8 @@ def jit_rope_then_store(
     out_loc: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.jit_kernel.kvcache import store_cache
     from sglang.jit_kernel.rope import apply_rope_inplace
+    from sglang.kernels.ops.kvcache.kvcache import store_cache
 
     head_size = q.shape[-1]
     row_dim = k.shape[-2] * head_size

test/registered/jit/benchmark/bench_store_cache.py

@@ -6,7 +6,7 @@ from sglang.jit_kernel.benchmark.utils import (
     create_empty,
     create_random,
 )
-from sglang.jit_kernel.kvcache import store_cache
+from sglang.kernels.ops.kvcache.kvcache import store_cache
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_tp_qknorm.py

@@ -23,16 +23,16 @@ import torch
 import torch.distributed as dist
 
 import sglang.srt.distributed.parallel_state as ps
-from sglang.jit_kernel.all_reduce import (
+from sglang.jit_kernel.benchmark import marker
+from sglang.jit_kernel.benchmark.utils import multigpu_bench_main
+from sglang.kernels.jit.utils import cache_once, get_ci_test_range
+from sglang.kernels.ops.communication.all_reduce import (
     fused_parallel_qknorm,
     get_all_reduce_module,
     get_fused_parallel_qknorm_max_occupancy,
     get_fused_parallel_qknorm_module,
 )
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import multigpu_bench_main
-from sglang.jit_kernel.mp import register_comm_cleanup
-from sglang.kernels.jit.utils import cache_once, get_ci_test_range
+from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )

test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py

@@ -4,7 +4,7 @@ per-buffer index_put_ stores (main K, main V, index K, optional index V)."""
 import torch
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
+from sglang.kernels.ops.kvcache.minimax_store_kv_index import store_kv_index
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/minimax/test_minimax_quant_scatter.py

@@ -4,11 +4,11 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.minimax_quant_ue8m0 import (
+from sglang.kernels.ops.moe.ep_moe_kernels import fill_gateup_input_triton_kernel
+from sglang.kernels.ops.quantization.minimax_quant_ue8m0 import (
     per_token_quant_fp8_ue8m0,
     per_token_quant_fp8_ue8m0_scatter,
 )
-from sglang.kernels.ops.moe.ep_moe_kernels import fill_gateup_input_triton_kernel
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")

test/registered/jit/minimax/test_minimax_store_kv_index.py

@@ -8,7 +8,7 @@ index_put_ stores would, for both value modes and int32/int64 indices.
 import pytest
 import torch
 
-from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
+from sglang.kernels.ops.kvcache.minimax_store_kv_index import store_kv_index
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_awq_dequantize.py

@@ -4,7 +4,9 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
+from sglang.kernels.ops.quantization.awq_dequantize import (
+    awq_dequantize as jit_awq_dequantize,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_awq_marlin_moe_repack.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 from sgl_kernel.scalar_type import scalar_types
 
-from sglang.jit_kernel.awq_marlin_repack import (
+from sglang.kernels.ops.quantization.awq_marlin_repack import (
     awq_marlin_moe_repack as jit_awq_marlin_moe_repack,
 )
 from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights

test/registered/jit/test_awq_marlin_repack.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 from sgl_kernel.scalar_type import scalar_types
 
-from sglang.jit_kernel.awq_marlin_repack import (
+from sglang.kernels.ops.quantization.awq_marlin_repack import (
     awq_marlin_repack as jit_awq_marlin_repack,
 )
 from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights

test/registered/jit/test_custom_all_reduce.py

@@ -28,10 +28,13 @@ import torch
 import torch.distributed as dist
 
 import sglang.srt.distributed.parallel_state as ps
-from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_all_reduce_module
-from sglang.jit_kernel.mp import register_comm_cleanup
 from sglang.jit_kernel.tests.utils import multigpu_pytest_main
 from sglang.kernels.jit.utils import cache_once, get_ci_test_range
+from sglang.kernels.ops.communication.all_reduce import (
+    AllReduceAlgo,
+    get_all_reduce_module,
+)
+from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )

test/registered/jit/test_cutedsl_bf16_gemm.py

@@ -17,7 +17,7 @@ register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b
 if not torch.cuda.is_available():
     pytest.skip("CUDA required", allow_module_level=True)
 
-from sglang.jit_kernel.cutedsl_bf16_gemm import cutedsl_bf16_gemm  # noqa: E402
+from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import cutedsl_bf16_gemm  # noqa: E402
 
 N_VALUES = [1024, 2624, 6144]
 K_VALUES = [2048, 6144]

test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py

@@ -17,7 +17,9 @@ register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-l
 if not torch.cuda.is_available():
     pytest.skip("CUDA required", allow_module_level=True)
 
-from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import dsv3_fused_a_gemm  # noqa: E402
+from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import (  # noqa: E402
+    dsv3_fused_a_gemm,
+)
 
 # hd_in must be a multiple of 256; 6144/7168 cover the real fused-A shapes.
 HD_INS = [6144, 7168]

test/registered/jit/test_fp8_blockwise_gemm.py

@@ -4,7 +4,7 @@ from typing import Optional, Type
 import pytest
 import torch
 
-from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
+from sglang.kernels.ops.gemm.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
 from sglang.srt.utils import is_sm120_supported
 from sglang.test.ci.ci_register import register_cuda_ci
 

test/registered/jit/test_fused_eh_norm.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.fused_eh_norm import fused_eh_norm
+from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_fused_fp8_qkv_kv_cache.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from sglang.jit_kernel.fused_fp8_qkv_kv_cache import fused_fp8_qkv_kv_cache
+from sglang.kernels.ops.kvcache.fused_fp8_qkv_kv_cache import fused_fp8_qkv_kv_cache
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_gptq_marlin.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 from sgl_kernel.scalar_type import scalar_types
 
-from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
+from sglang.kernels.ops.quantization.gptq_marlin import gptq_marlin_gemm
 from sglang.srt.layers.quantization.marlin_utils import (
     check_marlin_supported,
     marlin_make_workspace,

test/registered/jit/test_gptq_marlin_repack.py

@@ -4,7 +4,7 @@ import pytest
 import torch
 from sgl_kernel.scalar_type import scalar_types
 
-from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
 from sglang.srt.layers.quantization.utils import (
     gptq_quantize_weights,
     pack_rows,

test/registered/jit/test_hicache.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel
+from sglang.kernels.ops.kvcache.hicache import can_use_write_back_jit_kernel
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
 from sglang.srt.mem_cache.pool_host.common import (
     ALLOC_MEMORY_FUNCS,

test/registered/jit/test_hicache_page_first_write_back.py

@@ -13,7 +13,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel
+from sglang.kernels.ops.kvcache.hicache import can_use_write_back_jit_kernel
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
 from sglang.srt.mem_cache.pool_host.common import (
     ALLOC_MEMORY_FUNCS,

test/registered/jit/test_hisparse.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.hisparse import (
+from sglang.kernels.ops.kvcache.hisparse import (
     load_cache_to_device_buffer_dsv4_mla,
     load_cache_to_device_buffer_mla,
     transfer_cache_dsv4_mla,

test/registered/jit/test_mxfp8_moe.py

@@ -4,7 +4,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.mxfp8 import (
+from sglang.kernels.ops.quantization.mxfp8 import (
     es_sm100_mxfp8_blockscaled_grouped_quant,
     es_sm100_mxfp8_blockscaled_moe_grouped_gemm,
 )

test/registered/jit/test_ngram_embedding.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.ngram_embedding import (
+from sglang.kernels.ops.speculative.ngram_embedding import (
     compute_n_gram_ids,
     compute_n_gram_ids_decode,
     update_token_table,

test/registered/jit/test_resolve_future_token_ids.py

@@ -3,7 +3,9 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda
+from sglang.kernels.ops.speculative.resolve_future_token_ids import (
+    resolve_future_token_ids_cuda,
+)
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_rmsnorm_hf.py

@@ -6,11 +6,11 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.rmsnorm_hf import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.layernorm.rmsnorm_hf import (
     is_supported_rmsnorm_hf_hidden_size,
     rmsnorm_hf,
 )
-from sglang.kernels.jit.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_store_cache.py

@@ -4,8 +4,8 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.kvcache.kvcache import can_use_store_cache, store_cache
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=28, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_tp_qknorm.py

@@ -14,14 +14,14 @@ import torch.distributed as dist
 import triton
 
 import sglang.srt.distributed.parallel_state as ps
-from sglang.jit_kernel.all_reduce import (
+from sglang.jit_kernel.tests.utils import multigpu_pytest_main
+from sglang.kernels.jit.utils import cache_once
+from sglang.kernels.ops.communication.all_reduce import (
     fused_parallel_qknorm,
     get_all_reduce_module,
     get_fused_parallel_qknorm_module,
 )
-from sglang.jit_kernel.mp import register_comm_cleanup
-from sglang.jit_kernel.tests.utils import multigpu_pytest_main
-from sglang.kernels.jit.utils import cache_once
+from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )