diff --git a/python/sglang/jit_kernel/benchmark/utils.py b/python/sglang/jit_kernel/benchmark/utils.py index 74fed6b79..9f85a1945 100644 --- a/python/sglang/jit_kernel/benchmark/utils.py +++ b/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 diff --git a/python/sglang/jit_kernel/dsv4/attn.py b/python/sglang/jit_kernel/dsv4/attn.py index a436967a8..6f5b95841 100644 --- a/python/sglang/jit_kernel/dsv4/attn.py +++ b/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: diff --git a/python/sglang/jit_kernel/tests/utils.py b/python/sglang/jit_kernel/tests/utils.py index 1ead20913..6e6a1654f 100644 --- a/python/sglang/jit_kernel/tests/utils.py +++ b/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( diff --git a/python/sglang/jit_kernel/all_reduce.py b/python/sglang/kernels/ops/communication/all_reduce.py similarity index 100% rename from python/sglang/jit_kernel/all_reduce.py rename to python/sglang/kernels/ops/communication/all_reduce.py diff --git a/python/sglang/jit_kernel/mp.py b/python/sglang/kernels/ops/communication/mp.py similarity index 100% rename from python/sglang/jit_kernel/mp.py rename to python/sglang/kernels/ops/communication/mp.py diff --git a/python/sglang/jit_kernel/cutedsl_bf16_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py similarity index 100% rename from python/sglang/jit_kernel/cutedsl_bf16_gemm.py rename to python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py diff --git a/python/sglang/jit_kernel/cutedsl_dsv3_fused_a_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py similarity index 100% rename from python/sglang/jit_kernel/cutedsl_dsv3_fused_a_gemm.py rename to python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py diff --git a/python/sglang/jit_kernel/fp8_blockwise_gemm.py b/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py similarity index 100% rename from python/sglang/jit_kernel/fp8_blockwise_gemm.py rename to python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py diff --git a/python/sglang/jit_kernel/fused_a_gemm.py b/python/sglang/kernels/ops/gemm/fused_a_gemm.py similarity index 93% rename from python/sglang/jit_kernel/fused_a_gemm.py rename to python/sglang/kernels/ops/gemm/fused_a_gemm.py index 70b69bb49..72f012b32 100644 --- a/python/sglang/jit_kernel/fused_a_gemm.py +++ b/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) diff --git a/python/sglang/jit_kernel/fused_fp8_qkv_kv_cache.py b/python/sglang/kernels/ops/kvcache/fused_fp8_qkv_kv_cache.py similarity index 100% rename from python/sglang/jit_kernel/fused_fp8_qkv_kv_cache.py rename to python/sglang/kernels/ops/kvcache/fused_fp8_qkv_kv_cache.py diff --git a/python/sglang/jit_kernel/hicache.py b/python/sglang/kernels/ops/kvcache/hicache.py similarity index 100% rename from python/sglang/jit_kernel/hicache.py rename to python/sglang/kernels/ops/kvcache/hicache.py diff --git a/python/sglang/jit_kernel/hisparse.py b/python/sglang/kernels/ops/kvcache/hisparse.py similarity index 100% rename from python/sglang/jit_kernel/hisparse.py rename to python/sglang/kernels/ops/kvcache/hisparse.py diff --git a/python/sglang/jit_kernel/kvcache.py b/python/sglang/kernels/ops/kvcache/kvcache.py similarity index 100% rename from python/sglang/jit_kernel/kvcache.py rename to python/sglang/kernels/ops/kvcache/kvcache.py diff --git a/python/sglang/jit_kernel/minimax_store_kv_index.py b/python/sglang/kernels/ops/kvcache/minimax_store_kv_index.py similarity index 100% rename from python/sglang/jit_kernel/minimax_store_kv_index.py rename to python/sglang/kernels/ops/kvcache/minimax_store_kv_index.py diff --git a/python/sglang/jit_kernel/triton_store_cache.py b/python/sglang/kernels/ops/kvcache/triton_store_cache.py similarity index 100% rename from python/sglang/jit_kernel/triton_store_cache.py rename to python/sglang/kernels/ops/kvcache/triton_store_cache.py diff --git a/python/sglang/jit_kernel/fused_eh_norm.py b/python/sglang/kernels/ops/layernorm/fused_eh_norm.py similarity index 100% rename from python/sglang/jit_kernel/fused_eh_norm.py rename to python/sglang/kernels/ops/layernorm/fused_eh_norm.py diff --git a/python/sglang/jit_kernel/rmsnorm_hf.py b/python/sglang/kernels/ops/layernorm/rmsnorm_hf.py similarity index 100% rename from python/sglang/jit_kernel/rmsnorm_hf.py rename to python/sglang/kernels/ops/layernorm/rmsnorm_hf.py diff --git a/python/sglang/jit_kernel/inkling_sconv.py b/python/sglang/kernels/ops/mamba/inkling_sconv.py similarity index 100% rename from python/sglang/jit_kernel/inkling_sconv.py rename to python/sglang/kernels/ops/mamba/inkling_sconv.py diff --git a/python/sglang/jit_kernel/transfer_mamba.py b/python/sglang/kernels/ops/mamba/transfer_mamba.py similarity index 100% rename from python/sglang/jit_kernel/transfer_mamba.py rename to python/sglang/kernels/ops/mamba/transfer_mamba.py diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 8fb29be59..ed2ae1667 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/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, ) diff --git a/python/sglang/jit_kernel/awq_dequantize.py b/python/sglang/kernels/ops/quantization/awq_dequantize.py similarity index 100% rename from python/sglang/jit_kernel/awq_dequantize.py rename to python/sglang/kernels/ops/quantization/awq_dequantize.py diff --git a/python/sglang/jit_kernel/awq_marlin_repack.py b/python/sglang/kernels/ops/quantization/awq_marlin_repack.py similarity index 100% rename from python/sglang/jit_kernel/awq_marlin_repack.py rename to python/sglang/kernels/ops/quantization/awq_marlin_repack.py diff --git a/python/sglang/jit_kernel/fp8_quantize.py b/python/sglang/kernels/ops/quantization/fp8_quantize.py similarity index 100% rename from python/sglang/jit_kernel/fp8_quantize.py rename to python/sglang/kernels/ops/quantization/fp8_quantize.py diff --git a/python/sglang/jit_kernel/gptq_marlin.py b/python/sglang/kernels/ops/quantization/gptq_marlin.py similarity index 100% rename from python/sglang/jit_kernel/gptq_marlin.py rename to python/sglang/kernels/ops/quantization/gptq_marlin.py diff --git a/python/sglang/jit_kernel/gptq_marlin_repack.py b/python/sglang/kernels/ops/quantization/gptq_marlin_repack.py similarity index 100% rename from python/sglang/jit_kernel/gptq_marlin_repack.py rename to python/sglang/kernels/ops/quantization/gptq_marlin_repack.py diff --git a/python/sglang/jit_kernel/minimax_quant_ue8m0.py b/python/sglang/kernels/ops/quantization/minimax_quant_ue8m0.py similarity index 100% rename from python/sglang/jit_kernel/minimax_quant_ue8m0.py rename to python/sglang/kernels/ops/quantization/minimax_quant_ue8m0.py diff --git a/python/sglang/jit_kernel/mxfp8.py b/python/sglang/kernels/ops/quantization/mxfp8.py similarity index 100% rename from python/sglang/jit_kernel/mxfp8.py rename to python/sglang/kernels/ops/quantization/mxfp8.py diff --git a/python/sglang/jit_kernel/ngram_corpus.py b/python/sglang/kernels/ops/speculative/ngram_corpus.py similarity index 100% rename from python/sglang/jit_kernel/ngram_corpus.py rename to python/sglang/kernels/ops/speculative/ngram_corpus.py diff --git a/python/sglang/jit_kernel/ngram_embedding.py b/python/sglang/kernels/ops/speculative/ngram_embedding.py similarity index 100% rename from python/sglang/jit_kernel/ngram_embedding.py rename to python/sglang/kernels/ops/speculative/ngram_embedding.py diff --git a/python/sglang/jit_kernel/resolve_future_token_ids.py b/python/sglang/kernels/ops/speculative/resolve_future_token_ids.py similarity index 100% rename from python/sglang/jit_kernel/resolve_future_token_ids.py rename to python/sglang/kernels/ops/speculative/resolve_future_token_ids.py diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 031688419..047926280 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/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, diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py index 84386447f..91bffe663 100644 --- a/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py +++ b/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, ) diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py index 966e934c2..e3a0c8c52 100644 --- a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index 30f68f407..60826eeef 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index cefed569a..f32b59c12 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/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) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index aa151b720..be8e6377c 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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: diff --git a/python/sglang/srt/layers/n_gram_embedding.py b/python/sglang/srt/layers/n_gram_embedding.py index 5ec61f25d..d6420a792 100644 --- a/python/sglang/srt/layers/n_gram_embedding.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py index 16ef78bfb..ff3c7e16e 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py +++ b/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() diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index f8686a382..cc75aa730 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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") diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index 134a897cf..ffa210ff1 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/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__) diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py index f1463ae23..4ade14f15 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py +++ b/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() diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp8.py b/python/sglang/srt/layers/quantization/marlin_utils_fp8.py index 80afc12d4..6db4535da 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils_fp8.py +++ b/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() diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 5c77f22a7..282b1ced1 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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, ) diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 68c9d918d..3120ce0a3 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/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, ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 0140f9daa..f0879929e 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 498539cda..6b7ae4fcd 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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, ) diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 4569ce529..1af4bb972 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index 943395a50..bbfe420b5 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py b/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py index 71e8a3322..7b9e4347a 100644 --- a/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 4776ea19c..da3e63b86 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index a3395aed2..7114cd6d8 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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, ) diff --git a/python/sglang/srt/models/inkling_common/kernels/sconv.py b/python/sglang/srt/models/inkling_common/kernels/sconv.py index a9a790c58..2314aa1d9 100644 --- a/python/sglang/srt/models/inkling_common/kernels/sconv.py +++ b/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, ) diff --git a/python/sglang/srt/models/inkling_common/sconv.py b/python/sglang/srt/models/inkling_common/sconv.py index 8aa9435c7..5a793941e 100644 --- a/python/sglang/srt/models/inkling_common/sconv.py +++ b/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, ) diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 7fc27d1cd..abd228bfe 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/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 diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 476603ad3..00a405d4a 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py index e5a4b5656..100353339 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py +++ b/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__) diff --git a/test/manual/layers/moe/bench_mxfp4_sm90_kernels.py b/test/manual/layers/moe/bench_mxfp4_sm90_kernels.py index afbbcde34..b835612f7 100644 --- a/test/manual/layers/moe/bench_mxfp4_sm90_kernels.py +++ b/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, diff --git a/test/registered/jit/benchmark/bench_awq_dequantize.py b/test/registered/jit/benchmark/bench_awq_dequantize.py index 91b9d2dff..bd3842acc 100644 --- a/test/registered/jit/benchmark/bench_awq_dequantize.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_custom_all_reduce.py b/test/registered/jit/benchmark/bench_custom_all_reduce.py index b2ea3e9d6..667aa4375 100644 --- a/test/registered/jit/benchmark/bench_custom_all_reduce.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py index f66b89f27..429737b54 100644 --- a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py b/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py index fb603bfa0..514c155c2 100644 --- a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_fused_eh_norm.py b/test/registered/jit/benchmark/bench_fused_eh_norm.py index e000efc28..baff79839 100644 --- a/test/registered/jit/benchmark/bench_fused_eh_norm.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py b/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py index 9db0cebcc..24d4e5350 100644 --- a/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_hicache.py b/test/registered/jit/benchmark/bench_hicache.py index 10f261ac3..9bd45e535 100644 --- a/test/registered/jit/benchmark/bench_hicache.py +++ b/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, diff --git a/test/registered/jit/benchmark/bench_hisparse.py b/test/registered/jit/benchmark/bench_hisparse.py index 6d99729d0..cb589ce57 100644 --- a/test/registered/jit/benchmark/bench_hisparse.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_mxfp8_moe.py b/test/registered/jit/benchmark/bench_mxfp8_moe.py index c7c697208..d00fa1142 100644 --- a/test/registered/jit/benchmark/bench_mxfp8_moe.py +++ b/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, ) diff --git a/test/registered/jit/benchmark/bench_ngram_compute_decode.py b/test/registered/jit/benchmark/bench_ngram_compute_decode.py index 8801ab0a6..934c51063 100644 --- a/test/registered/jit/benchmark/bench_ngram_compute_decode.py +++ b/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, ) diff --git a/test/registered/jit/benchmark/bench_ngram_update_token_table.py b/test/registered/jit/benchmark/bench_ngram_update_token_table.py index 324f1c1c5..3198713b6 100644 --- a/test/registered/jit/benchmark/bench_ngram_update_token_table.py +++ b/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, ) diff --git a/test/registered/jit/benchmark/bench_resolve_future_token_ids.py b/test/registered/jit/benchmark/bench_resolve_future_token_ids.py index 8a88b5fb4..06c51c2d4 100644 --- a/test/registered/jit/benchmark/bench_resolve_future_token_ids.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_rope.py b/test/registered/jit/benchmark/bench_rope.py index 8e6ce7a0a..bc4bf3ec5 100644 --- a/test/registered/jit/benchmark/bench_rope.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_store_cache.py b/test/registered/jit/benchmark/bench_store_cache.py index 4df15acee..d70ba5f26 100644 --- a/test/registered/jit/benchmark/bench_store_cache.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_tp_qknorm.py b/test/registered/jit/benchmark/bench_tp_qknorm.py index 1f32defbe..02673ed86 100644 --- a/test/registered/jit/benchmark/bench_tp_qknorm.py +++ b/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, ) diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py b/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py index 5989147b1..29632bcc0 100644 --- a/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py +++ b/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( diff --git a/test/registered/jit/minimax/test_minimax_quant_scatter.py b/test/registered/jit/minimax/test_minimax_quant_scatter.py index 222aaf2dc..aaead0331 100644 --- a/test/registered/jit/minimax/test_minimax_quant_scatter.py +++ b/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") diff --git a/test/registered/jit/minimax/test_minimax_store_kv_index.py b/test/registered/jit/minimax/test_minimax_store_kv_index.py index 60902df35..ae94a6b17 100644 --- a/test/registered/jit/minimax/test_minimax_store_kv_index.py +++ b/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") diff --git a/test/registered/jit/test_awq_dequantize.py b/test/registered/jit/test_awq_dequantize.py index 8f3ba88c8..f0c47c826 100644 --- a/test/registered/jit/test_awq_dequantize.py +++ b/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") diff --git a/test/registered/jit/test_awq_marlin_moe_repack.py b/test/registered/jit/test_awq_marlin_moe_repack.py index caa191822..6c71e43ad 100644 --- a/test/registered/jit/test_awq_marlin_moe_repack.py +++ b/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 diff --git a/test/registered/jit/test_awq_marlin_repack.py b/test/registered/jit/test_awq_marlin_repack.py index d067ce6ef..bc3a8526e 100644 --- a/test/registered/jit/test_awq_marlin_repack.py +++ b/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 diff --git a/test/registered/jit/test_custom_all_reduce.py b/test/registered/jit/test_custom_all_reduce.py index 84571d81c..38ad20e95 100644 --- a/test/registered/jit/test_custom_all_reduce.py +++ b/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, ) diff --git a/test/registered/jit/test_cutedsl_bf16_gemm.py b/test/registered/jit/test_cutedsl_bf16_gemm.py index efb555cc6..3db73910c 100644 --- a/test/registered/jit/test_cutedsl_bf16_gemm.py +++ b/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] diff --git a/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py b/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py index 6fa9c224d..96a6004e4 100644 --- a/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py +++ b/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] diff --git a/test/registered/jit/test_fp8_blockwise_gemm.py b/test/registered/jit/test_fp8_blockwise_gemm.py index 5d097b9d2..be507ad9b 100644 --- a/test/registered/jit/test_fp8_blockwise_gemm.py +++ b/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 diff --git a/test/registered/jit/test_fused_eh_norm.py b/test/registered/jit/test_fused_eh_norm.py index 7149b6be7..b7f0ded07 100644 --- a/test/registered/jit/test_fused_eh_norm.py +++ b/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") diff --git a/test/registered/jit/test_fused_fp8_qkv_kv_cache.py b/test/registered/jit/test_fused_fp8_qkv_kv_cache.py index 2525b2666..c4e134a79 100644 --- a/test/registered/jit/test_fused_fp8_qkv_kv_cache.py +++ b/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") diff --git a/test/registered/jit/test_gptq_marlin.py b/test/registered/jit/test_gptq_marlin.py index 62e7eb7be..83c28c605 100644 --- a/test/registered/jit/test_gptq_marlin.py +++ b/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, diff --git a/test/registered/jit/test_gptq_marlin_repack.py b/test/registered/jit/test_gptq_marlin_repack.py index 870098c4a..6f633e0da 100644 --- a/test/registered/jit/test_gptq_marlin_repack.py +++ b/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, diff --git a/test/registered/jit/test_hicache.py b/test/registered/jit/test_hicache.py index 2b4d2d359..ec69f336b 100644 --- a/test/registered/jit/test_hicache.py +++ b/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, diff --git a/test/registered/jit/test_hicache_page_first_write_back.py b/test/registered/jit/test_hicache_page_first_write_back.py index 765e4817c..6b4f4870f 100644 --- a/test/registered/jit/test_hicache_page_first_write_back.py +++ b/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, diff --git a/test/registered/jit/test_hisparse.py b/test/registered/jit/test_hisparse.py index 7c80edeab..7f3104cce 100644 --- a/test/registered/jit/test_hisparse.py +++ b/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, diff --git a/test/registered/jit/test_mxfp8_moe.py b/test/registered/jit/test_mxfp8_moe.py index 1487f6748..ce2e17ad7 100644 --- a/test/registered/jit/test_mxfp8_moe.py +++ b/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, ) diff --git a/test/registered/jit/test_ngram_embedding.py b/test/registered/jit/test_ngram_embedding.py index 3a9dd2487..152c1fa1c 100644 --- a/test/registered/jit/test_ngram_embedding.py +++ b/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, diff --git a/test/registered/jit/test_resolve_future_token_ids.py b/test/registered/jit/test_resolve_future_token_ids.py index 51cde2724..6df314484 100644 --- a/test/registered/jit/test_resolve_future_token_ids.py +++ b/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") diff --git a/test/registered/jit/test_rmsnorm_hf.py b/test/registered/jit/test_rmsnorm_hf.py index 800b31671..ce00a5cbe 100644 --- a/test/registered/jit/test_rmsnorm_hf.py +++ b/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") diff --git a/test/registered/jit/test_store_cache.py b/test/registered/jit/test_store_cache.py index fdb620207..a114797d1 100644 --- a/test/registered/jit/test_store_cache.py +++ b/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") diff --git a/test/registered/jit/test_tp_qknorm.py b/test/registered/jit/test_tp_qknorm.py index ab0730f27..aa1434ea6 100644 --- a/test/registered/jit/test_tp_qknorm.py +++ b/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, )