[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>

python/sglang/jit_kernel/benchmark/utils.py

@@ -5,7 +5,7 @@ from typing import Callable, List, Optional, Sequence, Tuple
 import torch
 import triton.testing
 
-from sglang.jit_kernel.mp import multigpu_launch
+from sglang.kernels.ops.communication.mp import multigpu_launch
 from sglang.utils import is_in_ci
 
 

python/sglang/jit_kernel/dsv4/attn.py

@@ -60,7 +60,9 @@ def fused_store_cache(
     type: Literal["flashmla", "indexer"],
 ) -> None:
     if is_hip_runtime():
-        from sglang.jit_kernel.triton_store_cache import triton_fused_store_cache
+        from sglang.kernels.ops.kvcache.triton_store_cache import (
+            triton_fused_store_cache,
+        )
 
         triton_fused_store_cache(input, cache, indices, page_size=page_size, type=type)
     else:

python/sglang/jit_kernel/tests/utils.py

@@ -3,7 +3,7 @@ from typing import Callable, List, Optional, Sequence
 
 import pytest
 
-from sglang.jit_kernel.mp import multigpu_launch
+from sglang.kernels.ops.communication.mp import multigpu_launch
 
 
 def multigpu_pytest_main(

python/sglang/kernels/ops/gemm/fused_a_gemm.py

@@ -3,7 +3,7 @@
 Dispatches to one of two interchangeable implementations via ``backend``:
 
 - ``"jit"``: runtime-compiled CUDA C++ (``sglang.jit_kernel.dsv3_fused_a_gemm``).
-- ``"cutedsl"``: CuTe DSL (``sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm``).
+- ``"cutedsl"``: CuTe DSL (``sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm``).
 - ``"auto"``: CuTe DSL on SM120+, otherwise the JIT kernel.
 
 All backends share the signature ``(mat_a, mat_b, output=None) -> Tensor`` with
@@ -71,7 +71,7 @@ def dsv3_fused_a_gemm(
     if backend == FusedAGemmBackend.JIT:
         from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm as impl
     else:
-        from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import (
+        from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import (
             dsv3_fused_a_gemm as impl,
         )
     return impl(mat_a, mat_b, output)

python/sglang/kernels/ops/moe/ep_moe_kernels.py

@@ -1555,7 +1555,7 @@ def moe_ep_deepgemm_preprocess(
     block_n, block_k = block_shape[0], block_shape[1]
     is_fp8 = output_dtype == torch.float8_e4m3fn
     if is_fp8 and use_mxfp8:
-        from sglang.jit_kernel.minimax_quant_ue8m0 import (
+        from sglang.kernels.ops.quantization.minimax_quant_ue8m0 import (
             per_token_quant_fp8_ue8m0_scatter,
         )
 

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -25,7 +25,7 @@ import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
-from sglang.jit_kernel.all_reduce import (
+from sglang.kernels.ops.communication.all_reduce import (
     AllReduceAlgo,
     Communicator,
     IPCManager,

python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py

@@ -49,8 +49,8 @@ elif is_hip():
         pass
 else:
     try:
-        from sglang.jit_kernel.awq_dequantize import awq_dequantize
-        from sglang.jit_kernel.awq_marlin_repack import (
+        from sglang.kernels.ops.quantization.awq_dequantize import awq_dequantize
+        from sglang.kernels.ops.quantization.awq_marlin_repack import (
             awq_marlin_moe_repack,
             awq_marlin_repack,
         )

python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py

@@ -47,7 +47,7 @@ gptq_shuffle = _unsupported_kernel
 try:
     from sgl_kernel import gptq_gemm, gptq_shuffle
 
-    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+    from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
 except Exception:
     pass
 

python/sglang/srt/layers/attention/tokenspeed_mla_backend.py

@@ -33,9 +33,9 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
-from sglang.jit_kernel.fp8_quantize import fp8_quantize
 from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import mla_kv_pack_quantize_fp8
 from sglang.kernels.jit.utils import is_arch_support_pdl
+from sglang.kernels.ops.quantization.fp8_quantize import fp8_quantize
 from sglang.srt.layers.attention.trtllm_mla_backend import (
     TRTLLMMLABackend,
     TRTLLMMLAMultiStepDraftBackend,

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -664,7 +664,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         layer: RadixAttention,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor | None:
-        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,
+        )
 
         cache_loc = self._get_layer_cache_loc(layer, forward_batch)
         k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)

python/sglang/srt/layers/layernorm.py

@@ -130,8 +130,12 @@ if _is_cuda:
     # BEFORE the weight multiply, so the multiply is done in the narrow dtype.
     _jit_rmsnorm_hf_available = False
     try:
-        from sglang.jit_kernel.rmsnorm_hf import is_supported_rmsnorm_hf_hidden_size
-        from sglang.jit_kernel.rmsnorm_hf import rmsnorm_hf as _jit_rmsnorm_hf
+        from sglang.kernels.ops.layernorm.rmsnorm_hf import (
+            is_supported_rmsnorm_hf_hidden_size,
+        )
+        from sglang.kernels.ops.layernorm.rmsnorm_hf import (
+            rmsnorm_hf as _jit_rmsnorm_hf,
+        )
 
         _jit_rmsnorm_hf_available = True
     except ImportError:

python/sglang/srt/layers/n_gram_embedding.py

@@ -2,7 +2,7 @@ import torch
 from torch import nn
 from torch.nn import Parameter
 
-from sglang.jit_kernel.ngram_embedding import compute_n_gram_ids
+from sglang.kernels.ops.speculative.ngram_embedding import compute_n_gram_ids
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py

@@ -44,7 +44,7 @@ from sglang.srt.utils import is_cuda
 _is_cuda = is_cuda()
 
 if _is_cuda:
-    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+    from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
 
 
 ScalarType, scalar_types = get_scalar_types()

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -162,7 +162,7 @@ if _use_aiter:
 if _is_cuda:
     from sgl_kernel import fp8_scaled_mm
 
-    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.patch_torch import register_fake_if_exists
 
     @register_fake_if_exists("sgl_kernel::fp8_scaled_mm")

python/sglang/srt/layers/quantization/marlin_utils.py

@@ -46,7 +46,7 @@ except ImportError:
 _is_cuda = is_cuda()
 
 if _is_cuda:
-    from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
+    from sglang.kernels.ops.quantization.gptq_marlin import gptq_marlin_gemm
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/layers/quantization/marlin_utils_fp4.py

@@ -16,8 +16,8 @@ from sglang.srt.utils.custom_op import register_custom_op
 _is_cuda = is_cuda()
 
 if _is_cuda:
-    from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
-    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+    from sglang.kernels.ops.quantization.gptq_marlin import gptq_marlin_gemm
+    from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
 
 ScalarType, scalar_types = get_scalar_types()
 

python/sglang/srt/layers/quantization/marlin_utils_fp8.py

@@ -18,8 +18,8 @@ from sglang.srt.utils.custom_op import register_custom_op
 
 _is_cuda = is_cuda()
 if _is_cuda:
-    from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
-    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+    from sglang.kernels.ops.quantization.gptq_marlin import gptq_marlin_gemm
+    from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
 
 ScalarType, scalar_types = get_scalar_types()
 

python/sglang/srt/layers/quantization/unquant.py

@@ -93,7 +93,7 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
                 "--bf16-gemm-backend cutedsl requires SM100/SM103 (Blackwell)"
             )
 
-        from sglang.jit_kernel.cutedsl_bf16_gemm import (
+        from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import (
             cutedsl_bf16_gemm,
             use_cutedsl_bf16_gemm,
         )

python/sglang/srt/managers/hisparse_coordinator.py

@@ -5,7 +5,7 @@ from typing import List, NamedTuple, Union
 
 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,
 )

python/sglang/srt/mem_cache/memory_pool.py

@@ -38,7 +38,6 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
 from sglang.kernels.ops.attention.dsa import index_buf_accessor
 from sglang.kernels.ops.attention.dsa.quant_k_cache import (
     quantize_k_cache,
@@ -49,6 +48,7 @@ from sglang.kernels.ops.kvcache.cache_move import (
     set_kv_buffer_prefix_valid_tiled,
     store_cache_4d,
 )
+from sglang.kernels.ops.kvcache.kvcache import can_use_store_cache, store_cache
 from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
 from sglang.srt.configs.mamba_utils import BaseLinearStateParams
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
@@ -4902,7 +4902,7 @@ class MiniMaxSparseKVPool(KVCache):
         if index_pool is not None and self._can_fuse_kv_index_store(
             index_pool, cache_k, cache_idx_k
         ):
-            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
 
             main = self.main_pool
             head_bytes = main.head_dim * main.dtype.itemsize

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -13,13 +13,13 @@ import numpy as np
 import psutil
 import torch
 
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     can_use_write_back_jit_kernel,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
 )
-from sglang.jit_kernel.hisparse import transfer_cache_dsv4_mla
+from sglang.kernels.ops.kvcache.hisparse import transfer_cache_dsv4_mla
 from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, MambaPool
 from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
 
@@ -39,7 +39,7 @@ if _is_cuda or _is_hip:
         transfer_kv_per_layer_mla_pf_lf,
     )
 if _is_cuda:
-    from sglang.jit_kernel.transfer_mamba import (
+    from sglang.kernels.ops.mamba.transfer_mamba import (
         transfer_kv_mamba_lf_pf,
         transfer_kv_mamba_pf_lf,
     )

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

@@ -6,26 +6,26 @@ import threading
 import psutil
 import torch
 
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     can_use_hicache_jit_kernel,
     can_use_write_back_jit_kernel,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer as jit_transfer_hicache_all_layer,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer_mla as jit_transfer_hicache_all_layer_mla,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer_staged_lf_pf as jit_transfer_hicache_all_layer_staged_lf_pf,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_one_layer as jit_transfer_hicache_one_layer,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
 )
 from sglang.srt.mem_cache.memory_pool import MHATokenToKOnlyPool, MHATokenToKVPool

python/sglang/srt/mem_cache/pool_host/mla.py

@@ -5,17 +5,17 @@ from typing import Optional
 
 import torch
 
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     can_use_hicache_jit_kernel,
     can_use_write_back_jit_kernel,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer_mla as jit_transfer_hicache_all_layer_mla,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
 )
-from sglang.jit_kernel.hicache import (
+from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
 )
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool

python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
-from sglang.jit_kernel.ngram_embedding import update_token_table
+from sglang.kernels.ops.speculative.ngram_embedding import update_token_table
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.managers.schedule_batch import ForwardMode
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool

python/sglang/srt/models/deepseek_nextn.py

@@ -24,7 +24,7 @@ from safetensors.torch import load_file
 from torch import nn
 from transformers import PretrainedConfig
 
-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.srt.configs.model_config import is_deepseek_dsa
 from sglang.srt.distributed import get_pp_group
 from sglang.srt.environ import envs

python/sglang/srt/models/deepseek_v2.py

@@ -224,7 +224,7 @@ elif _is_npu:
 else:
     pass
 
-from sglang.jit_kernel.fused_a_gemm import (
+from sglang.kernels.ops.gemm.fused_a_gemm import (
     fused_a_gemm_weight_eligible,
     linear_with_fused_a_gemm,
 )

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

@@ -540,7 +540,9 @@ def causal_conv1d(
         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,
+        )
 
         return _cuda_causal_conv1d(
             x,
@@ -715,7 +717,7 @@ def update_sconv_cache(
         and x.stride(-1) == 1
         and sconv_cache.stride(2) == 1
     ):
-        from sglang.jit_kernel.inkling_sconv import (
+        from sglang.kernels.ops.mamba.inkling_sconv import (
             update_sconv_cache as _cuda_update_sconv_cache,
         )
 
@@ -970,7 +972,7 @@ def fused_causal_conv1d_update_decode(
         and x.stride(1) == 1
         and sconv_cache.stride(2) == 1
     ):
-        from sglang.jit_kernel.inkling_sconv import (
+        from sglang.kernels.ops.mamba.inkling_sconv import (
             fused_causal_conv1d_update_decode as _cuda_fused_decode,
         )
 

python/sglang/srt/models/inkling_common/sconv.py

@@ -814,7 +814,7 @@ def fused_gather_scatter_to_sconv_cache(
         and hidden_states.stride(-1) == 1
         and sconv_cache.stride(2) == 1
     ):
-        from sglang.jit_kernel.inkling_sconv import (
+        from sglang.kernels.ops.mamba.inkling_sconv import (
             fused_gather_scatter_to_sconv_cache as _cuda_gs,
         )
 
@@ -988,7 +988,7 @@ def fused_draft_extend_sconv_cache(
         and hidden_states.stride(-1) == 1
         and sconv_cache.stride(2) == 1
     ):
-        from sglang.jit_kernel.inkling_sconv import (
+        from sglang.kernels.ops.mamba.inkling_sconv import (
             fused_draft_extend_sconv_cache as _cuda_de,
         )
 

python/sglang/srt/models/minimax_m2.py

@@ -26,11 +26,11 @@ import triton.language as tl
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.jit_kernel.all_reduce import (
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.kernels.ops.communication.all_reduce import (
     fused_parallel_qknorm,
     get_fused_parallel_qknorm_max_occupancy,
 )
-from sglang.kernel_api_logging import debug_kernel_api
 from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
 from sglang.srt.distributed import get_pp_group, tensor_model_parallel_all_reduce
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder

python/sglang/srt/models/nemotron_h.py

@@ -103,7 +103,7 @@ from sglang.utils import logger
 _is_cuda = is_cuda()
 
 if _is_cuda:
-    from sglang.jit_kernel.fused_a_gemm import (
+    from sglang.kernels.ops.gemm.fused_a_gemm import (
         fused_a_gemm_weight_eligible,
         linear_with_fused_a_gemm,
     )

python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py

@@ -6,7 +6,7 @@ from typing import Dict, List, Tuple
 
 import numpy as np
 
-from sglang.jit_kernel.ngram_corpus import get_ngram_corpus_cls
+from sglang.kernels.ops.speculative.ngram_corpus import get_ngram_corpus_cls
 
 logger = logging.getLogger(__name__)