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

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,
 )