diff --git a/python/sglang/jit_kernel/__main__.py b/python/sglang/jit_kernel/__main__.py index b9c0f9681be5..0df3563c89e9 100644 --- a/python/sglang/jit_kernel/__main__.py +++ b/python/sglang/jit_kernel/__main__.py @@ -7,10 +7,10 @@ from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path -from sglang.jit_kernel.utils import get_jit_cuda_arch, override_jit_cuda_arch -from sglang.jit_kernel.utils.arch import get_default_target_flags -from sglang.jit_kernel.utils.compile import DEFAULT_INCLUDE -from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES +from sglang.kernels.jit import get_jit_cuda_arch, override_jit_cuda_arch +from sglang.kernels.jit.arch import get_default_target_flags +from sglang.kernels.jit.compile import DEFAULT_INCLUDE +from sglang.kernels.jit.deps import REGISTERED_DEPENDENCIES def _clangd_major_version() -> int | None: diff --git a/python/sglang/jit_kernel/activation.py b/python/sglang/jit_kernel/activation.py index e94aecb4c8a8..1e4c60eb1c7c 100644 --- a/python/sglang/jit_kernel/activation.py +++ b/python/sglang/jit_kernel/activation.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, get_jit_cuda_arch, is_arch_support_pdl, diff --git a/python/sglang/jit_kernel/add_constant.py b/python/sglang/jit_kernel/add_constant.py index 228e0de60d10..f3d042c66bf3 100644 --- a/python/sglang/jit_kernel/add_constant.py +++ b/python/sglang/jit_kernel/add_constant.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/all_reduce.py b/python/sglang/jit_kernel/all_reduce.py index ec8dd12f43d1..a9e5508b1f1d 100644 --- a/python/sglang/jit_kernel/all_reduce.py +++ b/python/sglang/jit_kernel/all_reduce.py @@ -7,14 +7,14 @@ import tvm_ffi from tvm_ffi import Module -from sglang.jit_kernel.utils import ( +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, lazy_register_class, load_jit, make_cpp_args, ) -from sglang.kernel_api_logging import debug_kernel_api class AllReduceAlgo(enum.Enum): diff --git a/python/sglang/jit_kernel/awq_dequantize.py b/python/sglang/jit_kernel/awq_dequantize.py index 4a188c02e51b..e2399a44be44 100644 --- a/python/sglang/jit_kernel/awq_dequantize.py +++ b/python/sglang/jit_kernel/awq_dequantize.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/awq_marlin_repack.py b/python/sglang/jit_kernel/awq_marlin_repack.py index d51c1fd5194d..ce5ffda65338 100644 --- a/python/sglang/jit_kernel/awq_marlin_repack.py +++ b/python/sglang/jit_kernel/awq_marlin_repack.py @@ -4,8 +4,8 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/benchmark/marker.py b/python/sglang/jit_kernel/benchmark/marker.py index 60e31783b8ed..486ef5ecbcfc 100644 --- a/python/sglang/jit_kernel/benchmark/marker.py +++ b/python/sglang/jit_kernel/benchmark/marker.py @@ -21,7 +21,7 @@ import torch -from sglang.jit_kernel.utils import cache_once +from sglang.kernels.jit import cache_once from sglang.utils import is_in_ci F = TypeVar("F", bound=Callable[..., "BenchResult"]) diff --git a/python/sglang/jit_kernel/clamp_position.py b/python/sglang/jit_kernel/clamp_position.py index ed57da776666..f934457876db 100644 --- a/python/sglang/jit_kernel/clamp_position.py +++ b/python/sglang/jit_kernel/clamp_position.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/concat_mla.py b/python/sglang/jit_kernel/concat_mla.py index 4945b73bc27f..74153734aa41 100644 --- a/python/sglang/jit_kernel/concat_mla.py +++ b/python/sglang/jit_kernel/concat_mla.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py b/python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py index 60ff5626d2d8..1d2c24a854b2 100644 --- a/python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py +++ b/python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py b/python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py index 45277e2d2cee..24bd61b1db36 100644 --- a/python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py +++ b/python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/diffusion/norm_scale_shift_native.py b/python/sglang/jit_kernel/diffusion/norm_scale_shift_native.py index cbc175918e2a..1f00a957dbdd 100644 --- a/python/sglang/jit_kernel/diffusion/norm_scale_shift_native.py +++ b/python/sglang/jit_kernel/diffusion/norm_scale_shift_native.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/diffusion/qknorm_rope.py b/python/sglang/jit_kernel/diffusion/qknorm_rope.py index 8dfdf8d8db21..1901199b006e 100644 --- a/python/sglang/jit_kernel/diffusion/qknorm_rope.py +++ b/python/sglang/jit_kernel/diffusion/qknorm_rope.py @@ -5,7 +5,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/diffusion/residual_gate_add.py b/python/sglang/jit_kernel/diffusion/residual_gate_add.py index 9933e3d2aa09..de1e1ba7c3cd 100644 --- a/python/sglang/jit_kernel/diffusion/residual_gate_add.py +++ b/python/sglang/jit_kernel/diffusion/residual_gate_add.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/dsv32/elementwise.py b/python/sglang/jit_kernel/dsv32/elementwise.py index 4c07afe4b77f..fada8630302c 100644 --- a/python/sglang/jit_kernel/dsv32/elementwise.py +++ b/python/sglang/jit_kernel/dsv32/elementwise.py @@ -2,7 +2,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/dsv3_fused_a_gemm.py b/python/sglang/jit_kernel/dsv3_fused_a_gemm.py index c718a5131925..16bd4992f0d5 100644 --- a/python/sglang/jit_kernel/dsv3_fused_a_gemm.py +++ b/python/sglang/jit_kernel/dsv3_fused_a_gemm.py @@ -11,13 +11,13 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) -from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.common import direct_register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/dsv3_router_gemm.py b/python/sglang/jit_kernel/dsv3_router_gemm.py index f33d883ddbd3..1a564d5eeb63 100644 --- a/python/sglang/jit_kernel/dsv3_router_gemm.py +++ b/python/sglang/jit_kernel/dsv3_router_gemm.py @@ -11,13 +11,13 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) -from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/dsv4/attn.py b/python/sglang/jit_kernel/dsv4/attn.py index 784711498b27..3a02da9477ed 100644 --- a/python/sglang/jit_kernel/dsv4/attn.py +++ b/python/sglang/jit_kernel/dsv4/attn.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, is_hip_runtime, diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index bf82bff04214..72b824060bf3 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/python/sglang/jit_kernel/dsv4/compress.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/dsv4/compress_old.py b/python/sglang/jit_kernel/dsv4/compress_old.py index 9bf96a964aab..de0c1f261b00 100644 --- a/python/sglang/jit_kernel/dsv4/compress_old.py +++ b/python/sglang/jit_kernel/dsv4/compress_old.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/dsv4/elementwise.py b/python/sglang/jit_kernel/dsv4/elementwise.py index b1395824ddbc..15acb355dd95 100644 --- a/python/sglang/jit_kernel/dsv4/elementwise.py +++ b/python/sglang/jit_kernel/dsv4/elementwise.py @@ -2,7 +2,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/dsv4/fp8_wo_a.py b/python/sglang/jit_kernel/dsv4/fp8_wo_a.py index 907900bdb4cb..c174dde498c4 100644 --- a/python/sglang/jit_kernel/dsv4/fp8_wo_a.py +++ b/python/sglang/jit_kernel/dsv4/fp8_wo_a.py @@ -4,13 +4,13 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) -from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.custom_op import register_custom_op from .utils import make_name diff --git a/python/sglang/jit_kernel/dsv4/moe.py b/python/sglang/jit_kernel/dsv4/moe.py index 81a21c85c837..5aa85c05e5a7 100644 --- a/python/sglang/jit_kernel/dsv4/moe.py +++ b/python/sglang/jit_kernel/dsv4/moe.py @@ -2,7 +2,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, is_hip_runtime, diff --git a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py b/python/sglang/jit_kernel/dsv4/online_c128_mtp.py index 5a901b4d9bd8..a0dfd6bacae0 100644 --- a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py +++ b/python/sglang/jit_kernel/dsv4/online_c128_mtp.py @@ -6,7 +6,7 @@ import torch from sglang.jit_kernel.dsv4.utils import make_name -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args from sglang.srt.environ import envs if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/dsv4/topk.py b/python/sglang/jit_kernel/dsv4/topk.py index fdaedfd75685..329fcab5ac31 100644 --- a/python/sglang/jit_kernel/dsv4/topk.py +++ b/python/sglang/jit_kernel/dsv4/topk.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, is_hip_runtime, diff --git a/python/sglang/jit_kernel/fixup_zero_kv.py b/python/sglang/jit_kernel/fixup_zero_kv.py index 6175c0f37dc9..dfa653464477 100644 --- a/python/sglang/jit_kernel/fixup_zero_kv.py +++ b/python/sglang/jit_kernel/fixup_zero_kv.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index fe7f42234b17..cc3966e7da17 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/python/sglang/jit_kernel/flash_attention_v3.py @@ -4,8 +4,8 @@ import torch -from sglang.jit_kernel.utils import cache_once from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once from sglang.srt.environ import envs from sglang.srt.utils import get_device_capability, is_musa diff --git a/python/sglang/jit_kernel/flash_attn/cute/interface.py b/python/sglang/jit_kernel/flash_attn/cute/interface.py index cc8ae4e21bad..a631423d309f 100644 --- a/python/sglang/jit_kernel/flash_attn/cute/interface.py +++ b/python/sglang/jit_kernel/flash_attn/cute/interface.py @@ -15,7 +15,7 @@ from sglang.jit_kernel.flash_attn.cute.cache_utils import get_jit_cache from sglang.jit_kernel.flash_attn.cute.testing import is_fake_mode -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None: from sglang.jit_kernel.flash_attn.cute import cute_dsl_ptxas # noqa: F401 diff --git a/python/sglang/jit_kernel/fp8_blockwise_gemm.py b/python/sglang/jit_kernel/fp8_blockwise_gemm.py index 49b4c9606e25..065adf34f2ed 100644 --- a/python/sglang/jit_kernel/fp8_blockwise_gemm.py +++ b/python/sglang/jit_kernel/fp8_blockwise_gemm.py @@ -5,8 +5,8 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit, override_jit_cuda_arch from sglang.srt.utils.common import is_sm120_supported from sglang.srt.utils.custom_op import register_custom_op diff --git a/python/sglang/jit_kernel/fused_eh_norm.py b/python/sglang/jit_kernel/fused_eh_norm.py index d7b7c0817ebc..e90943de93f4 100644 --- a/python/sglang/jit_kernel/fused_eh_norm.py +++ b/python/sglang/jit_kernel/fused_eh_norm.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/fused_fp8_qkv_kv_cache.py b/python/sglang/jit_kernel/fused_fp8_qkv_kv_cache.py index d412cae5594f..6142058cacf8 100644 --- a/python/sglang/jit_kernel/fused_fp8_qkv_kv_cache.py +++ b/python/sglang/jit_kernel/fused_fp8_qkv_kv_cache.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/fused_metadata_copy.py b/python/sglang/jit_kernel/fused_metadata_copy.py index 68d0f9227812..fc752891692a 100644 --- a/python/sglang/jit_kernel/fused_metadata_copy.py +++ b/python/sglang/jit_kernel/fused_metadata_copy.py @@ -15,7 +15,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args logger = logging.getLogger(__name__) diff --git a/python/sglang/jit_kernel/fused_qknorm_rope.py b/python/sglang/jit_kernel/fused_qknorm_rope.py index 00e872020709..ed5a5060d028 100644 --- a/python/sglang/jit_kernel/fused_qknorm_rope.py +++ b/python/sglang/jit_kernel/fused_qknorm_rope.py @@ -5,7 +5,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/fused_store_index_cache.py b/python/sglang/jit_kernel/fused_store_index_cache.py index f8b3b1432183..5c50ba8a0a1f 100644 --- a/python/sglang/jit_kernel/fused_store_index_cache.py +++ b/python/sglang/jit_kernel/fused_store_index_cache.py @@ -13,13 +13,13 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) -from sglang.kernel_api_logging import debug_kernel_api if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/gptq_marlin.py b/python/sglang/jit_kernel/gptq_marlin.py index d3bde5336476..198fea661881 100644 --- a/python/sglang/jit_kernel/gptq_marlin.py +++ b/python/sglang/jit_kernel/gptq_marlin.py @@ -4,8 +4,8 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from sgl_kernel.scalar_type import ScalarType diff --git a/python/sglang/jit_kernel/gptq_marlin_repack.py b/python/sglang/jit_kernel/gptq_marlin_repack.py index ea7fe9908b9c..e3f740fc47a1 100644 --- a/python/sglang/jit_kernel/gptq_marlin_repack.py +++ b/python/sglang/jit_kernel/gptq_marlin_repack.py @@ -4,8 +4,8 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/hadamard.py b/python/sglang/jit_kernel/hadamard.py index 6e845474903e..4ebf0199c7fd 100644 --- a/python/sglang/jit_kernel/hadamard.py +++ b/python/sglang/jit_kernel/hadamard.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import KERNEL_PATH, cache_once, load_jit, make_cpp_args from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/hicache.py b/python/sglang/jit_kernel/hicache.py index 268eabe18948..d430568fca33 100644 --- a/python/sglang/jit_kernel/hicache.py +++ b/python/sglang/jit_kernel/hicache.py @@ -3,8 +3,8 @@ import logging from typing import TYPE_CHECKING -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: import torch diff --git a/python/sglang/jit_kernel/hisparse.py b/python/sglang/jit_kernel/hisparse.py index b88143da1b44..8258ad0eee68 100644 --- a/python/sglang/jit_kernel/hisparse.py +++ b/python/sglang/jit_kernel/hisparse.py @@ -5,7 +5,7 @@ import torch -from sglang.jit_kernel.utils import load_jit, make_cpp_args +from sglang.kernels.jit import load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/inkling_all_reduce.py b/python/sglang/jit_kernel/inkling_all_reduce.py index b8dcbca3ed68..5fd84aaa9259 100644 --- a/python/sglang/jit_kernel/inkling_all_reduce.py +++ b/python/sglang/jit_kernel/inkling_all_reduce.py @@ -10,7 +10,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, empty_sentinel, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, empty_sentinel, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/inkling_ar_fused.py b/python/sglang/jit_kernel/inkling_ar_fused.py index 455bf5379b3b..d078dc05a15d 100644 --- a/python/sglang/jit_kernel/inkling_ar_fused.py +++ b/python/sglang/jit_kernel/inkling_ar_fused.py @@ -9,7 +9,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, empty_sentinel, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, empty_sentinel, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/inkling_ar_scattered_sconv.py b/python/sglang/jit_kernel/inkling_ar_scattered_sconv.py index ed742e139a7a..683a18b0891a 100644 --- a/python/sglang/jit_kernel/inkling_ar_scattered_sconv.py +++ b/python/sglang/jit_kernel/inkling_ar_scattered_sconv.py @@ -10,7 +10,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/inkling_attn_prologue.py b/python/sglang/jit_kernel/inkling_attn_prologue.py index 50a00b3e5fd4..e63cafbe2c82 100644 --- a/python/sglang/jit_kernel/inkling_attn_prologue.py +++ b/python/sglang/jit_kernel/inkling_attn_prologue.py @@ -7,7 +7,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, empty_sentinel, is_arch_support_pdl, diff --git a/python/sglang/jit_kernel/inkling_gate_topk_renorm.py b/python/sglang/jit_kernel/inkling_gate_topk_renorm.py index 46ccf054c86b..19e11b7814a8 100644 --- a/python/sglang/jit_kernel/inkling_gate_topk_renorm.py +++ b/python/sglang/jit_kernel/inkling_gate_topk_renorm.py @@ -22,7 +22,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/inkling_rel_proj.py b/python/sglang/jit_kernel/inkling_rel_proj.py index ebefd9cfd980..2ed183ba382b 100644 --- a/python/sglang/jit_kernel/inkling_rel_proj.py +++ b/python/sglang/jit_kernel/inkling_rel_proj.py @@ -10,7 +10,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, empty_sentinel, is_arch_support_pdl, diff --git a/python/sglang/jit_kernel/inkling_row_scale.py b/python/sglang/jit_kernel/inkling_row_scale.py index 4b831410ae63..13ede858daac 100644 --- a/python/sglang/jit_kernel/inkling_row_scale.py +++ b/python/sglang/jit_kernel/inkling_row_scale.py @@ -9,7 +9,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/inkling_sconv.py b/python/sglang/jit_kernel/inkling_sconv.py index a3121dceae05..21b498d82c86 100644 --- a/python/sglang/jit_kernel/inkling_sconv.py +++ b/python/sglang/jit_kernel/inkling_sconv.py @@ -10,7 +10,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/kpool_topk_transform.py b/python/sglang/jit_kernel/kpool_topk_transform.py index 52dfdcf61bf9..da6a5498b73f 100644 --- a/python/sglang/jit_kernel/kpool_topk_transform.py +++ b/python/sglang/jit_kernel/kpool_topk_transform.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py b/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py index f63de2484090..19468017e9fc 100644 --- a/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py +++ b/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/kv_canary/verify.py b/python/sglang/jit_kernel/kv_canary/verify.py index b7fd260bd24d..c2b755cc912a 100644 --- a/python/sglang/jit_kernel/kv_canary/verify.py +++ b/python/sglang/jit_kernel/kv_canary/verify.py @@ -7,7 +7,7 @@ import torch from sglang.jit_kernel.kv_canary import consts -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/kv_canary/write.py b/python/sglang/jit_kernel/kv_canary/write.py index 6926a2dc90e1..21c3383a082d 100644 --- a/python/sglang/jit_kernel/kv_canary/write.py +++ b/python/sglang/jit_kernel/kv_canary/write.py @@ -11,7 +11,7 @@ _assert_contiguous, _build_real_kv_source_abi, ) -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/kvcache.py b/python/sglang/jit_kernel/kvcache.py index b611e76570d1..1e96867a5d34 100644 --- a/python/sglang/jit_kernel/kvcache.py +++ b/python/sglang/jit_kernel/kvcache.py @@ -5,7 +5,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/lplb/cuda_solver.py b/python/sglang/jit_kernel/lplb/cuda_solver.py index 60a78c751b48..87977744279b 100644 --- a/python/sglang/jit_kernel/lplb/cuda_solver.py +++ b/python/sglang/jit_kernel/lplb/cuda_solver.py @@ -17,7 +17,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, get_jit_cuda_arch, load_jit, diff --git a/python/sglang/jit_kernel/minimax_decode_topk.py b/python/sglang/jit_kernel/minimax_decode_topk.py index 26940d66d8a3..cce646ac8716 100644 --- a/python/sglang/jit_kernel/minimax_decode_topk.py +++ b/python/sglang/jit_kernel/minimax_decode_topk.py @@ -17,7 +17,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/minimax_qknorm_rope.py b/python/sglang/jit_kernel/minimax_qknorm_rope.py index 7815c398f9b6..d24830aea19e 100644 --- a/python/sglang/jit_kernel/minimax_qknorm_rope.py +++ b/python/sglang/jit_kernel/minimax_qknorm_rope.py @@ -22,7 +22,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/minimax_quant_ue8m0.py b/python/sglang/jit_kernel/minimax_quant_ue8m0.py index 240e9cf75a64..8af9f5ad735e 100644 --- a/python/sglang/jit_kernel/minimax_quant_ue8m0.py +++ b/python/sglang/jit_kernel/minimax_quant_ue8m0.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/minimax_store_kv_index.py b/python/sglang/jit_kernel/minimax_store_kv_index.py index 2c5fd440815f..e82283d4b9c2 100644 --- a/python/sglang/jit_kernel/minimax_store_kv_index.py +++ b/python/sglang/jit_kernel/minimax_store_kv_index.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/mla_kv_pack_quantize_fp8.py b/python/sglang/jit_kernel/mla_kv_pack_quantize_fp8.py index bf1024eb09c5..aa3afd194828 100644 --- a/python/sglang/jit_kernel/mla_kv_pack_quantize_fp8.py +++ b/python/sglang/jit_kernel/mla_kv_pack_quantize_fp8.py @@ -11,7 +11,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl @triton.jit diff --git a/python/sglang/jit_kernel/moe_align.py b/python/sglang/jit_kernel/moe_align.py index ee136f1a8a46..381c3ad66cbb 100644 --- a/python/sglang/jit_kernel/moe_align.py +++ b/python/sglang/jit_kernel/moe_align.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/moe_finalize_fuse_shared.py b/python/sglang/jit_kernel/moe_finalize_fuse_shared.py index 7432397327db..71f13c1e44e5 100644 --- a/python/sglang/jit_kernel/moe_finalize_fuse_shared.py +++ b/python/sglang/jit_kernel/moe_finalize_fuse_shared.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit @cache_once diff --git a/python/sglang/jit_kernel/moe_fused_gate.py b/python/sglang/jit_kernel/moe_fused_gate.py index 6cdfd9c3a6b6..c45b6fee98e5 100644 --- a/python/sglang/jit_kernel/moe_fused_gate.py +++ b/python/sglang/jit_kernel/moe_fused_gate.py @@ -7,8 +7,8 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import cache_once, is_arch_support_pdl, load_jit from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, is_arch_support_pdl, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/moe_lora_align.py b/python/sglang/jit_kernel/moe_lora_align.py index f18ad7ca0771..260f6fd7840f 100644 --- a/python/sglang/jit_kernel/moe_lora_align.py +++ b/python/sglang/jit_kernel/moe_lora_align.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/moe_permute_prepare.py b/python/sglang/jit_kernel/moe_permute_prepare.py index 679c52332203..03870e98c069 100644 --- a/python/sglang/jit_kernel/moe_permute_prepare.py +++ b/python/sglang/jit_kernel/moe_permute_prepare.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/moe_topk_sigmoid.py b/python/sglang/jit_kernel/moe_topk_sigmoid.py index 7fbdd22a6014..160651daa20d 100644 --- a/python/sglang/jit_kernel/moe_topk_sigmoid.py +++ b/python/sglang/jit_kernel/moe_topk_sigmoid.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/moe_wna16_marlin.py b/python/sglang/jit_kernel/moe_wna16_marlin.py index e9a8cd25372b..0a7bf21bfd40 100644 --- a/python/sglang/jit_kernel/moe_wna16_marlin.py +++ b/python/sglang/jit_kernel/moe_wna16_marlin.py @@ -4,8 +4,8 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from sgl_kernel.scalar_type import ScalarType diff --git a/python/sglang/jit_kernel/mxfp8.py b/python/sglang/jit_kernel/mxfp8.py index 2f0a91f9f1ba..2771a4b4fc30 100644 --- a/python/sglang/jit_kernel/mxfp8.py +++ b/python/sglang/jit_kernel/mxfp8.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, load_jit, make_cpp_args, diff --git a/python/sglang/jit_kernel/ngram_corpus.py b/python/sglang/jit_kernel/ngram_corpus.py index d2121417c8d4..437eb3b839e8 100644 --- a/python/sglang/jit_kernel/ngram_corpus.py +++ b/python/sglang/jit_kernel/ngram_corpus.py @@ -7,7 +7,7 @@ import torch import tvm_ffi -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit _MATCH_TYPE_MAP = {"BFS": 0, "PROB": 1} diff --git a/python/sglang/jit_kernel/ngram_embedding.py b/python/sglang/jit_kernel/ngram_embedding.py index ea7da20ba15d..3a7a4875b3cb 100644 --- a/python/sglang/jit_kernel/ngram_embedding.py +++ b/python/sglang/jit_kernel/ngram_embedding.py @@ -2,8 +2,8 @@ from typing import TYPE_CHECKING -from sglang.jit_kernel.utils import cache_once, load_jit from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: import torch diff --git a/python/sglang/jit_kernel/norm.py b/python/sglang/jit_kernel/norm.py index 4fb3c9451b3a..8f8c14e46087 100644 --- a/python/sglang/jit_kernel/norm.py +++ b/python/sglang/jit_kernel/norm.py @@ -1,179 +1,11 @@ -from __future__ import annotations +"""Compatibility shim (RFC #29630 Phase 4). -import logging -from typing import TYPE_CHECKING, Optional +The JIT rmsnorm/qknorm operators moved to +``sglang.kernels.ops.layernorm._jit_norm``. This module re-exports them so +existing ``sglang.jit_kernel.norm`` import paths keep working; it will be +removed in Phase 5. +""" -import torch +from sglang.kernels.ops.layernorm import _jit_norm as _impl -from sglang.jit_kernel.utils import ( - cache_once, - is_arch_support_pdl, - load_jit, - make_cpp_args, -) -from sglang.kernel_api_logging import debug_kernel_api - -if TYPE_CHECKING: - from tvm_ffi.module import Module - - -logger = logging.getLogger(__name__) - - -@cache_once -def _jit_qknorm_module(head_dim: int, dtype: torch.dtype) -> Module: - args = make_cpp_args(head_dim, is_arch_support_pdl(), dtype) - return load_jit( - "qknorm", - *args, - cuda_files=["elementwise/qknorm.cuh"], - cuda_wrappers=[("qknorm", f"QKNormKernel<{args}>::run")], - ) - - -_RMSNORM_WARP_SIZES = frozenset({64, 128, 256}) -_RMSNORM_MAX_HIDDEN_SIZE = 16384 -_RMSNORM_HALF_BLOCK_MIN_SIZE = 2048 - - -def _is_supported_rmsnorm_hidden_size(d: int) -> bool: - return d in _RMSNORM_WARP_SIZES or ( - (d > 256 and d % 256 == 0 and d <= 8192) - or (d >= 8192 and d % 512 == 0 and d <= 16384) - ) - - -def _rmsnorm_kernel_class(hidden_size: int) -> str: - if hidden_size in _RMSNORM_WARP_SIZES: - return "RMSNormWarpKernel" - if hidden_size == 512: - return "RMSNormHalfKernel" - if hidden_size >= _RMSNORM_HALF_BLOCK_MIN_SIZE: - if hidden_size % 512 == 0: - return "RMSNormHalfKernel" - return "RMSNormKernel" - - -@cache_once -def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module: - args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype) - kernel_class = f"{_rmsnorm_kernel_class(hidden_size)}<{args}>" - return load_jit( - "rmsnorm", - *args, - cuda_files=["elementwise/rmsnorm.cuh"], - cuda_wrappers=[("rmsnorm", f"{kernel_class}::run")], - ) - - -def is_supported_jit_fused_add_rmsnorm_hidden_size(hidden_size: int) -> bool: - return hidden_size > 0 and hidden_size % 16 == 0 and hidden_size <= 8192 - - -@cache_once -def _jit_fused_add_rmsnorm_module( - dtype: torch.dtype, cast_x_before_out_mul: bool -) -> Module: - args = make_cpp_args(cast_x_before_out_mul, dtype) - return load_jit( - "fused_add_rmsnorm", - *args, - cuda_files=["elementwise/fused_add_rmsnorm.cuh"], - cuda_wrappers=[("fused_add_rmsnorm", f"FusedAddRMSNormKernel<{args}>::run")], - ) - - -@cache_once -def _jit_qknorm_across_heads_module(dtype: torch.dtype) -> Module: - args = make_cpp_args(dtype) - return load_jit( - "qknorm_across_heads", - *args, - cuda_files=["elementwise/qknorm_across_heads.cuh"], - cuda_wrappers=[ - ("qknorm_across_heads", f"QKNormAcrossHeadsKernel<{args}>::run") - ], - ) - - -@torch.compiler.assume_constant_result -@cache_once -def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool: - if head_dim not in [64, 128, 256, 512, 1024]: - logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel") - return False - try: - _jit_qknorm_module(head_dim, dtype) - return True - except Exception as e: - logger.warning(f"Failed to load JIT QK-Norm kernel: {e}") - return False - - -@debug_kernel_api -def fused_inplace_qknorm( - q: torch.Tensor, - k: torch.Tensor, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - eps: float = 1e-6, - *, - head_dim: int = 0, -) -> None: - head_dim = head_dim or q.size(-1) - module = _jit_qknorm_module(head_dim, q.dtype) - module.qknorm(q, k, q_weight, k_weight, eps) - - -@debug_kernel_api -def rmsnorm( - input: torch.Tensor, - weight: torch.Tensor, - out: Optional[torch.Tensor] = None, - eps: float = 1e-6, -) -> None: - out = out if out is not None else input - hidden_size = input.size(-1) - if not _is_supported_rmsnorm_hidden_size(hidden_size): - raise RuntimeError( - f"jit rmsnorm: unsupported hidden_size={hidden_size}. " - f"Supported: {sorted(_RMSNORM_WARP_SIZES)}, and multiples of 256 in " - f"(256, {_RMSNORM_MAX_HIDDEN_SIZE}]." - ) - module = _jit_rmsnorm_module(hidden_size, input.dtype) - module.rmsnorm(input, weight, out, eps) - - -@debug_kernel_api -def fused_add_rmsnorm( - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, - *, - cast_x_before_out_mul: bool = False, -) -> None: - module = _jit_fused_add_rmsnorm_module(input.dtype, cast_x_before_out_mul) - module.fused_add_rmsnorm(input, residual, weight, eps) - - -@debug_kernel_api -def fused_inplace_qknorm_across_heads( - q: torch.Tensor, - k: torch.Tensor, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - eps: float = 1e-6, -) -> None: - """ - Fused inplace QK normalization across all heads. - - Args: - q: Query tensor of shape [batch_size, num_heads * head_dim] - k: Key tensor of shape [batch_size, num_heads * head_dim] - q_weight: Query weight tensor of shape [num_heads * head_dim] - k_weight: Key weight tensor of shape [num_heads * head_dim] - eps: Epsilon for numerical stability - """ - module = _jit_qknorm_across_heads_module(q.dtype) - module.qknorm_across_heads(q, k, q_weight, k_weight, eps) +globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/per_tensor_quant_fp8.py b/python/sglang/jit_kernel/per_tensor_quant_fp8.py index 40aee2607cb8..12c13de6edd7 100644 --- a/python/sglang/jit_kernel/per_tensor_quant_fp8.py +++ b/python/sglang/jit_kernel/per_tensor_quant_fp8.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/per_token_group_quant_8bit.py b/python/sglang/jit_kernel/per_token_group_quant_8bit.py index cf4b753a89d9..7574358e200c 100644 --- a/python/sglang/jit_kernel/per_token_group_quant_8bit.py +++ b/python/sglang/jit_kernel/per_token_group_quant_8bit.py @@ -4,13 +4,13 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) -from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py b/python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py index 8c0ebbfc237b..75adea5f0593 100644 --- a/python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py +++ b/python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py @@ -4,13 +4,13 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) -from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/resolve_future_token_ids.py b/python/sglang/jit_kernel/resolve_future_token_ids.py index d0068a4ad7b6..71918f21dadb 100644 --- a/python/sglang/jit_kernel/resolve_future_token_ids.py +++ b/python/sglang/jit_kernel/resolve_future_token_ids.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/rmsnorm_hf.py b/python/sglang/jit_kernel/rmsnorm_hf.py index f4db56daccae..3062176877db 100644 --- a/python/sglang/jit_kernel/rmsnorm_hf.py +++ b/python/sglang/jit_kernel/rmsnorm_hf.py @@ -6,7 +6,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/rope.py b/python/sglang/jit_kernel/rope.py index d9cbe0a8b324..f2ed5ff50eee 100644 --- a/python/sglang/jit_kernel/rope.py +++ b/python/sglang/jit_kernel/rope.py @@ -5,7 +5,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/set_mla_kv_buffer.py b/python/sglang/jit_kernel/set_mla_kv_buffer.py index 3624f5afc324..d59b8cc81b28 100644 --- a/python/sglang/jit_kernel/set_mla_kv_buffer.py +++ b/python/sglang/jit_kernel/set_mla_kv_buffer.py @@ -13,7 +13,7 @@ import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit import ( cache_once, is_arch_support_pdl, load_jit, diff --git a/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py b/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py index 372f30c828c9..352c01f96c84 100644 --- a/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py +++ b/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py @@ -10,8 +10,8 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit, override_jit_cuda_arch from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/timestep_embedding.py b/python/sglang/jit_kernel/timestep_embedding.py index 08a941a6ee72..31a57c01d481 100644 --- a/python/sglang/jit_kernel/timestep_embedding.py +++ b/python/sglang/jit_kernel/timestep_embedding.py @@ -4,8 +4,8 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/transfer_mamba.py b/python/sglang/jit_kernel/transfer_mamba.py index a64afa06067f..f894160488ed 100644 --- a/python/sglang/jit_kernel/transfer_mamba.py +++ b/python/sglang/jit_kernel/transfer_mamba.py @@ -4,7 +4,7 @@ and ``transfer_kv_mamba_lf_pf`` (backup: layer_first -> page_first). Uses the shared ``load_jit`` + ``cache_once`` infrastructure from -``sglang.jit_kernel.utils`` — the same mechanism used by ``hicache.py`` +``sglang.kernels.jit`` — the same mechanism used by ``hicache.py`` for MHA/MLA staged write-back kernels. This ensures consistent content-addressed caching, CUDA arch detection, and multi-worker JIT compilation behavior across all JIT kernels. @@ -15,8 +15,8 @@ import logging from typing import TYPE_CHECKING -from sglang.jit_kernel.utils import cache_once, load_jit from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: import torch diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/kimi_k2_moe_fused_gate.py b/python/sglang/jit_kernel/trtllm_lora_temp/kimi_k2_moe_fused_gate.py index 0fd6231d3f58..936b5bcd77ac 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/kimi_k2_moe_fused_gate.py +++ b/python/sglang/jit_kernel/trtllm_lora_temp/kimi_k2_moe_fused_gate.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/moe_lora_merged_align.py b/python/sglang/jit_kernel/trtllm_lora_temp/moe_lora_merged_align.py index 32b3fa8b5212..fbe50a736ae6 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/moe_lora_merged_align.py +++ b/python/sglang/jit_kernel/trtllm_lora_temp/moe_lora_merged_align.py @@ -4,7 +4,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/topk_softmax_pack.py b/python/sglang/jit_kernel/trtllm_lora_temp/topk_softmax_pack.py index 36a4cdef7b5c..e91f837d9f50 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/topk_softmax_pack.py +++ b/python/sglang/jit_kernel/trtllm_lora_temp/topk_softmax_pack.py @@ -20,7 +20,7 @@ import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.kernels.jit import cache_once, load_jit from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: diff --git a/python/sglang/jit_kernel/utils/__init__.py b/python/sglang/kernels/jit/__init__.py similarity index 71% rename from python/sglang/jit_kernel/utils/__init__.py rename to python/sglang/kernels/jit/__init__.py index 452ee54410cc..1a9fcef9a478 100644 --- a/python/sglang/jit_kernel/utils/__init__.py +++ b/python/sglang/kernels/jit/__init__.py @@ -1,11 +1,11 @@ -"""Public interface of sglang.jit_kernel.utils.""" +"""Public interface of sglang.kernels.jit.""" -from sglang.jit_kernel.utils.arch import ( +from sglang.kernels.jit.arch import ( get_jit_cuda_arch, is_arch_support_pdl, override_jit_cuda_arch, ) -from sglang.jit_kernel.utils.common import ( +from sglang.kernels.jit.common import ( cache_once, empty_sentinel, get_ci_test_range, @@ -14,7 +14,7 @@ lazy_register_class, should_run_full_tests, ) -from sglang.jit_kernel.utils.compile import KERNEL_PATH, load_jit, make_cpp_args +from sglang.kernels.jit.compile import KERNEL_PATH, load_jit, make_cpp_args __all__ = [ "empty_sentinel", diff --git a/python/sglang/jit_kernel/utils/arch.py b/python/sglang/kernels/jit/arch.py similarity index 98% rename from python/sglang/jit_kernel/utils/arch.py rename to python/sglang/kernels/jit/arch.py index 24c6492eefa3..6d6dbfa83282 100644 --- a/python/sglang/jit_kernel/utils/arch.py +++ b/python/sglang/kernels/jit/arch.py @@ -9,7 +9,7 @@ import torch -from sglang.jit_kernel.utils.common import ( +from sglang.kernels.jit.common import ( cache_once, is_hip_runtime, is_musa_runtime, diff --git a/python/sglang/jit_kernel/utils/common.py b/python/sglang/kernels/jit/common.py similarity index 100% rename from python/sglang/jit_kernel/utils/common.py rename to python/sglang/kernels/jit/common.py diff --git a/python/sglang/jit_kernel/utils/compile.py b/python/sglang/kernels/jit/compile.py similarity index 98% rename from python/sglang/jit_kernel/utils/compile.py rename to python/sglang/kernels/jit/compile.py index 1196d8559891..a3d8db96e442 100644 --- a/python/sglang/jit_kernel/utils/compile.py +++ b/python/sglang/kernels/jit/compile.py @@ -13,9 +13,9 @@ import torch -from sglang.jit_kernel.utils.arch import get_default_target_flags, get_jit_cuda_arch -from sglang.jit_kernel.utils.common import cache_once, is_hip_runtime -from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES +from sglang.kernels.jit.arch import get_default_target_flags, get_jit_cuda_arch +from sglang.kernels.jit.common import cache_once, is_hip_runtime +from sglang.kernels.jit.deps import REGISTERED_DEPENDENCIES if TYPE_CHECKING: from tvm_ffi import Module diff --git a/python/sglang/jit_kernel/utils/deps.py b/python/sglang/kernels/jit/deps.py similarity index 100% rename from python/sglang/jit_kernel/utils/deps.py rename to python/sglang/kernels/jit/deps.py diff --git a/python/sglang/kernels/ops/attention/dsv4/compress_c128_hip.py b/python/sglang/kernels/ops/attention/dsv4/compress_c128_hip.py index 6838ebe8bafc..50914fbe9e30 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress_c128_hip.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress_c128_hip.py @@ -14,7 +14,7 @@ CompressorDecodePlan, CompressorPrefillPlan, ) -from sglang.jit_kernel.utils import is_hip_runtime +from sglang.kernels.jit import is_hip_runtime _is_hip = is_hip_runtime() diff --git a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py index 0fa7e032f785..3e0224ee6e4d 100644 --- a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py +++ b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py @@ -15,7 +15,7 @@ import triton.language as tl from einops import rearrange -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled from sglang.srt.model_executor.cuda_graph_config import ( Backend, diff --git a/python/sglang/kernels/ops/attention/utils.py b/python/sglang/kernels/ops/attention/utils.py index 74b1308031bd..8e3de9fe4c11 100644 --- a/python/sglang/kernels/ops/attention/utils.py +++ b/python/sglang/kernels/ops/attention/utils.py @@ -2,7 +2,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.kernels.ops.attention.pad import ( pad_sequence_with_mask as pad_sequence_with_mask, ) diff --git a/python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py index 541d1ccf6059..edeca80adb8e 100644 --- a/python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py +++ b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py @@ -1,7 +1,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl def get_pdl_launch_metadata() -> tuple[bool, dict]: diff --git a/python/sglang/kernels/ops/kvcache/mla_buffer.py b/python/sglang/kernels/ops/kvcache/mla_buffer.py index 5bf285d97965..b095778cdebc 100644 --- a/python/sglang/kernels/ops/kvcache/mla_buffer.py +++ b/python/sglang/kernels/ops/kvcache/mla_buffer.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.runtime_context import get_parallel diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index 9f9953114139..d42dd8560f56 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -113,7 +113,7 @@ def forward_jit( ) -> torch.Tensor: import torch - from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm + from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm if out is None: out = torch.empty_like(input) @@ -227,7 +227,9 @@ def forward_jit( eps: float = 1e-6, enable_pdl: Optional[bool] = None, ) -> None: - from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm + from sglang.kernels.ops.layernorm._jit_norm import ( + fused_add_rmsnorm as jit_fused_add_rmsnorm, + ) return jit_fused_add_rmsnorm(input, residual, weight, eps) diff --git a/python/sglang/kernels/ops/layernorm/_jit_norm.py b/python/sglang/kernels/ops/layernorm/_jit_norm.py new file mode 100644 index 000000000000..f601b915689d --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/_jit_norm.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.jit import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +logger = logging.getLogger(__name__) + + +@cache_once +def _jit_qknorm_module(head_dim: int, dtype: torch.dtype) -> Module: + args = make_cpp_args(head_dim, is_arch_support_pdl(), dtype) + return load_jit( + "qknorm", + *args, + cuda_files=["elementwise/qknorm.cuh"], + cuda_wrappers=[("qknorm", f"QKNormKernel<{args}>::run")], + ) + + +_RMSNORM_WARP_SIZES = frozenset({64, 128, 256}) +_RMSNORM_MAX_HIDDEN_SIZE = 16384 +_RMSNORM_HALF_BLOCK_MIN_SIZE = 2048 + + +def _is_supported_rmsnorm_hidden_size(d: int) -> bool: + return d in _RMSNORM_WARP_SIZES or ( + (d > 256 and d % 256 == 0 and d <= 8192) + or (d >= 8192 and d % 512 == 0 and d <= 16384) + ) + + +def _rmsnorm_kernel_class(hidden_size: int) -> str: + if hidden_size in _RMSNORM_WARP_SIZES: + return "RMSNormWarpKernel" + if hidden_size == 512: + return "RMSNormHalfKernel" + if hidden_size >= _RMSNORM_HALF_BLOCK_MIN_SIZE: + if hidden_size % 512 == 0: + return "RMSNormHalfKernel" + return "RMSNormKernel" + + +@cache_once +def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module: + args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype) + kernel_class = f"{_rmsnorm_kernel_class(hidden_size)}<{args}>" + return load_jit( + "rmsnorm", + *args, + cuda_files=["elementwise/rmsnorm.cuh"], + cuda_wrappers=[("rmsnorm", f"{kernel_class}::run")], + ) + + +def is_supported_jit_fused_add_rmsnorm_hidden_size(hidden_size: int) -> bool: + return hidden_size > 0 and hidden_size % 16 == 0 and hidden_size <= 8192 + + +@cache_once +def _jit_fused_add_rmsnorm_module( + dtype: torch.dtype, cast_x_before_out_mul: bool +) -> Module: + args = make_cpp_args(cast_x_before_out_mul, dtype) + return load_jit( + "fused_add_rmsnorm", + *args, + cuda_files=["elementwise/fused_add_rmsnorm.cuh"], + cuda_wrappers=[("fused_add_rmsnorm", f"FusedAddRMSNormKernel<{args}>::run")], + ) + + +@cache_once +def _jit_qknorm_across_heads_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + return load_jit( + "qknorm_across_heads", + *args, + cuda_files=["elementwise/qknorm_across_heads.cuh"], + cuda_wrappers=[ + ("qknorm_across_heads", f"QKNormAcrossHeadsKernel<{args}>::run") + ], + ) + + +@torch.compiler.assume_constant_result +@cache_once +def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool: + if head_dim not in [64, 128, 256, 512, 1024]: + logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel") + return False + try: + _jit_qknorm_module(head_dim, dtype) + return True + except Exception as e: + logger.warning(f"Failed to load JIT QK-Norm kernel: {e}") + return False + + +@debug_kernel_api +def fused_inplace_qknorm( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float = 1e-6, + *, + head_dim: int = 0, +) -> None: + head_dim = head_dim or q.size(-1) + module = _jit_qknorm_module(head_dim, q.dtype) + module.qknorm(q, k, q_weight, k_weight, eps) + + +@debug_kernel_api +def rmsnorm( + input: torch.Tensor, + weight: torch.Tensor, + out: Optional[torch.Tensor] = None, + eps: float = 1e-6, +) -> None: + out = out if out is not None else input + hidden_size = input.size(-1) + if not _is_supported_rmsnorm_hidden_size(hidden_size): + raise RuntimeError( + f"jit rmsnorm: unsupported hidden_size={hidden_size}. " + f"Supported: {sorted(_RMSNORM_WARP_SIZES)}, and multiples of 256 in " + f"(256, {_RMSNORM_MAX_HIDDEN_SIZE}]." + ) + module = _jit_rmsnorm_module(hidden_size, input.dtype) + module.rmsnorm(input, weight, out, eps) + + +@debug_kernel_api +def fused_add_rmsnorm( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + *, + cast_x_before_out_mul: bool = False, +) -> None: + module = _jit_fused_add_rmsnorm_module(input.dtype, cast_x_before_out_mul) + module.fused_add_rmsnorm(input, residual, weight, eps) + + +@debug_kernel_api +def fused_inplace_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float = 1e-6, +) -> None: + """ + Fused inplace QK normalization across all heads. + + Args: + q: Query tensor of shape [batch_size, num_heads * head_dim] + k: Key tensor of shape [batch_size, num_heads * head_dim] + q_weight: Query weight tensor of shape [num_heads * head_dim] + k_weight: Key weight tensor of shape [num_heads * head_dim] + eps: Epsilon for numerical stability + """ + module = _jit_qknorm_across_heads_module(q.dtype) + module.qknorm_across_heads(q, k, q_weight, k_weight, eps) diff --git a/python/sglang/kernels/ops/layernorm/elementwise.py b/python/sglang/kernels/ops/layernorm/elementwise.py index 1414e0038bb7..0a97bd17ac01 100644 --- a/python/sglang/kernels/ops/layernorm/elementwise.py +++ b/python/sglang/kernels/ops/layernorm/elementwise.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.kernels.ops.activation.softcap import softcap_out as fused_softcap from sglang.srt.utils import is_hip from sglang.srt.utils.custom_op import register_custom_op diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index 190d40874440..6c381d3c7546 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -7,7 +7,7 @@ import torch -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py index d70f176e86ef..badec109c705 100644 --- a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py +++ b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl PAD_SLOT_ID = -1 diff --git a/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py b/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py index 09dbd73db2b7..7127719290ed 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py @@ -11,7 +11,7 @@ import triton.language as tl from packaging import version -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl PAD_SLOT_ID = -1 diff --git a/python/sglang/kernels/ops/quantization/fp8_kernel.py b/python/sglang/kernels/ops/quantization/fp8_kernel.py index e6b5793ad313..450466a1f619 100644 --- a/python/sglang/kernels/ops/quantization/fp8_kernel.py +++ b/python/sglang/kernels/ops/quantization/fp8_kernel.py @@ -28,7 +28,7 @@ except: pass -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.utils import ( diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index 12f5a8bf288b..f52e9668a973 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -10,7 +10,7 @@ compress_forward, compress_norm_rope_store, ) -from sglang.jit_kernel.utils import is_hip_runtime +from sglang.kernels.jit import is_hip_runtime from sglang.srt.environ import envs if TYPE_CHECKING: diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index 75b4ae06a806..6436a3b75b8e 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -35,7 +35,7 @@ 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.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.layers.attention.trtllm_mla_backend import ( TRTLLMMLABackend, TRTLLMMLAMultiStepDraftBackend, diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 72d44a474465..d73a0b1d938b 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -75,7 +75,7 @@ def finalize_flashinfer_trtllm_deferred_output( shared_output: torch.Tensor, ) -> torch.Tensor: from sglang.jit_kernel.moe_finalize_fuse_shared import moe_finalize_fuse_shared - from sglang.jit_kernel.utils import is_arch_support_pdl + from sglang.kernels.jit import is_arch_support_pdl return moe_finalize_fuse_shared( deferred_output.gemm2_out, diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/inkling_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/inkling_moe.py index e7e3c2f28139..a6e1af7188a4 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/inkling_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/inkling_moe.py @@ -6,7 +6,7 @@ import triton import triton.language as tl -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.layers.moe.moe_runner.triton_utils.helion_utils import ( get_model_depths, helion_aot_autotune, diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py index bd3c3c7afab8..1c7e2df065eb 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py @@ -13,7 +13,7 @@ import triton.language as tl from sglang.jit_kernel.inkling_gate_topk_renorm import inkling_gate_topk_renorm_v2 -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.environ import envs from sglang.srt.layers.moe.moe_runner.triton_utils.gate_topk import ( fpval_to_key, diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index eae718dd1276..4fa904ab4b17 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/python/sglang/srt/models/gpt_oss.py @@ -26,7 +26,7 @@ from torch import nn from transformers import PretrainedConfig -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.distributed import ( get_pp_group, tensor_model_parallel_all_reduce, diff --git a/python/sglang/srt/models/inkling_common/moe.py b/python/sglang/srt/models/inkling_common/moe.py index 9746d9a5c49e..38d73c9c7748 100644 --- a/python/sglang/srt/models/inkling_common/moe.py +++ b/python/sglang/srt/models/inkling_common/moe.py @@ -14,7 +14,7 @@ inkling_gate_gemv, inkling_gate_gemv_fused, ) -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.configs.inkling import InklingModelConfig from sglang.srt.distributed import ( get_tensor_model_parallel_group, diff --git a/test/registered/jit/benchmark/bench_custom_all_reduce.py b/test/registered/jit/benchmark/bench_custom_all_reduce.py index 4bc187693138..a146b5d0d061 100644 --- a/test/registered/jit/benchmark/bench_custom_all_reduce.py +++ b/test/registered/jit/benchmark/bench_custom_all_reduce.py @@ -13,7 +13,7 @@ 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.jit_kernel.utils import cache_once, is_arch_support_pdl +from sglang.kernels.jit import cache_once, is_arch_support_pdl 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 ef084d79d3a5..6741a42d8be5 100644 --- a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py +++ b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py @@ -15,7 +15,7 @@ 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.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.jit import get_jit_cuda_arch, is_hip_runtime from sglang.srt.utils.common import is_sm120_supported 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_dsv3_router_gemm.py b/test/registered/jit/benchmark/bench_dsv3_router_gemm.py index 91018638c5f4..6188f3258755 100644 --- a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py +++ b/test/registered/jit/benchmark/bench_dsv3_router_gemm.py @@ -15,7 +15,7 @@ from sglang.jit_kernel.benchmark import marker from sglang.jit_kernel.benchmark.utils import create_random from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm -from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.jit import get_jit_cuda_arch, is_hip_runtime from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( diff --git a/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py b/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py index aa3f4080d88a..554acb02327d 100644 --- a/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py +++ b/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py @@ -17,7 +17,7 @@ from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import ( mla_kv_pack_quantize_fp8 as hybrid_pack, ) -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( diff --git a/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py b/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py index 28f05e964bdc..49d52a949741 100644 --- a/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py +++ b/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py @@ -21,7 +21,7 @@ get_benchmark_range, ) from sglang.jit_kernel.set_mla_kv_buffer import set_mla_kv_buffer as jit_set -from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.kernels.jit import is_arch_support_pdl from sglang.srt.mem_cache.utils import set_mla_kv_buffer_kernel as sglang_triton_kernel from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton as sglang_wrapper from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_symm_mem_all_gather.py b/test/registered/jit/benchmark/bench_symm_mem_all_gather.py index 86c31ff84b03..ae2cd4e8dab5 100644 --- a/test/registered/jit/benchmark/bench_symm_mem_all_gather.py +++ b/test/registered/jit/benchmark/bench_symm_mem_all_gather.py @@ -28,7 +28,7 @@ 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.utils import cache_once +from sglang.kernels.jit import cache_once from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( all_gather_inner, create_state, diff --git a/test/registered/jit/benchmark/bench_tp_qknorm.py b/test/registered/jit/benchmark/bench_tp_qknorm.py index 290f5b71c2f8..dad4b85f9910 100644 --- a/test/registered/jit/benchmark/bench_tp_qknorm.py +++ b/test/registered/jit/benchmark/bench_tp_qknorm.py @@ -32,7 +32,7 @@ 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.jit_kernel.utils import cache_once, get_ci_test_range +from sglang.kernels.jit import cache_once, get_ci_test_range from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) diff --git a/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py index cfabb8ff59de..9af718ddaddb 100644 --- a/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py +++ b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py @@ -11,7 +11,7 @@ import torch from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE -from sglang.jit_kernel.utils import KERNEL_PATH +from sglang.kernels.jit import KERNEL_PATH from sglang.test.ci.ci_register import register_cuda_ci from sglang.utils import is_in_ci diff --git a/test/registered/jit/benchmark/diffusion/bench_norm_impls.py b/test/registered/jit/benchmark/diffusion/bench_norm_impls.py index 489f1069393c..ec440b575d2a 100644 --- a/test/registered/jit/benchmark/diffusion/bench_norm_impls.py +++ b/test/registered/jit/benchmark/diffusion/bench_norm_impls.py @@ -18,7 +18,7 @@ from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm -from sglang.jit_kernel.utils import KERNEL_PATH +from sglang.kernels.jit import KERNEL_PATH from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.utils import is_in_ci diff --git a/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py b/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py index 25c2a3b0bf92..af605e030797 100644 --- a/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py +++ b/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py @@ -9,7 +9,7 @@ from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import ( fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_amd_ci, 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/diffusion/test_qknorm_rope.py b/test/registered/jit/diffusion/test_qknorm_rope.py index 8d775d7b825b..b2944e1c28b2 100644 --- a/test/registered/jit/diffusion/test_qknorm_rope.py +++ b/test/registered/jit/diffusion/test_qknorm_rope.py @@ -5,7 +5,7 @@ import torch import triton -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=44, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/diffusion/test_qwen_image_modulation.py b/test/registered/jit/diffusion/test_qwen_image_modulation.py index 596f6bcc7683..0e556e02bb7c 100644 --- a/test/registered/jit/diffusion/test_qwen_image_modulation.py +++ b/test/registered/jit/diffusion/test_qwen_image_modulation.py @@ -9,7 +9,7 @@ fuse_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/diffusion/test_varlen_pack_pad.py b/test/registered/jit/diffusion/test_varlen_pack_pad.py index 614c4b711e6d..1162e25251c3 100644 --- a/test/registered/jit/diffusion/test_varlen_pack_pad.py +++ b/test/registered/jit/diffusion/test_varlen_pack_pad.py @@ -12,7 +12,7 @@ fused_pack_qkv, fused_scatter_to_padded, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py b/test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py index 51837e17dec2..4bf925f40f89 100644 --- a/test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py +++ b/test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py @@ -20,7 +20,7 @@ fused_scatter_to_padded, ) from sglang.jit_kernel.flash_attention import flash_attn_varlen_func -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.multimodal_gen.runtime.layers.attention.backends import ( flash_attn as _fa_backend, ) diff --git a/test/registered/jit/test_activation.py b/test/registered/jit/test_activation.py index 62c2becba916..1eec5ff754da 100644 --- a/test/registered/jit/test_activation.py +++ b/test/registered/jit/test_activation.py @@ -9,7 +9,7 @@ relu2, run_activation, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_custom_all_reduce.py b/test/registered/jit/test_custom_all_reduce.py index 8fd238828a7a..01fd4a4ec42a 100644 --- a/test/registered/jit/test_custom_all_reduce.py +++ b/test/registered/jit/test_custom_all_reduce.py @@ -31,7 +31,7 @@ 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.jit_kernel.utils import cache_once, get_ci_test_range +from sglang.kernels.jit import cache_once, get_ci_test_range 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 9d293d5295ac..ddab7f5ba6e3 100644 --- a/test/registered/jit/test_cutedsl_bf16_gemm.py +++ b/test/registered/jit/test_cutedsl_bf16_gemm.py @@ -5,7 +5,7 @@ import pytest import torch -from sglang.jit_kernel.utils import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.jit import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") 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 96993fc01bfd..e283fad42962 100644 --- a/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py +++ b/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py @@ -5,7 +5,7 @@ import pytest import torch -from sglang.jit_kernel.utils import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.jit import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime from sglang.test.ci.ci_register import 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_dsv3_fused_a_gemm.py b/test/registered/jit/test_dsv3_fused_a_gemm.py index 83287f0e0f4b..34f4963de41d 100644 --- a/test/registered/jit/test_dsv3_fused_a_gemm.py +++ b/test/registered/jit/test_dsv3_fused_a_gemm.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm -from sglang.jit_kernel.utils import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.jit import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime from sglang.test.ci.ci_register import 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_dsv3_router_gemm.py b/test/registered/jit/test_dsv3_router_gemm.py index b934fe9030c2..a7df28643aff 100644 --- a/test/registered/jit/test_dsv3_router_gemm.py +++ b/test/registered/jit/test_dsv3_router_gemm.py @@ -7,7 +7,7 @@ import torch from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm -from sglang.jit_kernel.utils import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.jit import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=37, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_fused_add_rmsnorm.py b/test/registered/jit/test_fused_add_rmsnorm.py index 8edb56345100..30cebf15e42c 100644 --- a/test/registered/jit/test_fused_add_rmsnorm.py +++ b/test/registered/jit/test_fused_add_rmsnorm.py @@ -4,7 +4,7 @@ import pytest import torch -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_mla_kv_pack_quantize_fp8.py b/test/registered/jit/test_mla_kv_pack_quantize_fp8.py index 3266a4715ae5..69095c7c46ed 100644 --- a/test/registered/jit/test_mla_kv_pack_quantize_fp8.py +++ b/test/registered/jit/test_mla_kv_pack_quantize_fp8.py @@ -4,7 +4,7 @@ import torch from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import mla_kv_pack_quantize_fp8 -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_moe_align_block_size.py b/test/registered/jit/test_moe_align_block_size.py index 78bf6e121303..5f151a44b1fc 100644 --- a/test/registered/jit/test_moe_align_block_size.py +++ b/test/registered/jit/test_moe_align_block_size.py @@ -7,7 +7,7 @@ import triton.language as tl from sglang.jit_kernel.moe_align import moe_align_block_size -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import 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_moe_fused_gate.py b/test/registered/jit/test_moe_fused_gate.py index 07ac21f66350..e90c00275617 100644 --- a/test/registered/jit/test_moe_fused_gate.py +++ b/test/registered/jit/test_moe_fused_gate.py @@ -22,7 +22,7 @@ import torch from sglang.jit_kernel.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_per_tensor_quant_fp8.py b/test/registered/jit/test_per_tensor_quant_fp8.py index 75a389453fb7..f388d06feaf2 100644 --- a/test/registered/jit/test_per_tensor_quant_fp8.py +++ b/test/registered/jit/test_per_tensor_quant_fp8.py @@ -6,7 +6,7 @@ import torch from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_per_token_group_quant_8bit.py b/test/registered/jit/test_per_token_group_quant_8bit.py index 41b5d7dc576b..4b158b20558b 100644 --- a/test/registered/jit/test_per_token_group_quant_8bit.py +++ b/test/registered/jit/test_per_token_group_quant_8bit.py @@ -7,7 +7,7 @@ from sglang.jit_kernel.per_token_group_quant_8bit import ( per_token_group_quant_8bit as sglang_per_token_group_quant_8bit, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.srt.utils import is_hip from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/test_per_token_group_quant_8bit_v2.py b/test/registered/jit/test_per_token_group_quant_8bit_v2.py index e87adeb7eedb..139f671bbe24 100644 --- a/test/registered/jit/test_per_token_group_quant_8bit_v2.py +++ b/test/registered/jit/test_per_token_group_quant_8bit_v2.py @@ -6,7 +6,7 @@ from sglang.jit_kernel.per_token_group_quant_8bit_v2 import ( per_token_group_quant_8bit_v2, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_qknorm.py b/test/registered/jit/test_qknorm.py index a540ae3a1d87..832be7f8cf63 100644 --- a/test/registered/jit/test_qknorm.py +++ b/test/registered/jit/test_qknorm.py @@ -5,7 +5,7 @@ import torch import triton -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=37, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_qknorm_across_heads.py b/test/registered/jit/test_qknorm_across_heads.py index e21258b558f7..c956d7111eae 100644 --- a/test/registered/jit/test_qknorm_across_heads.py +++ b/test/registered/jit/test_qknorm_across_heads.py @@ -5,7 +5,7 @@ import torch import triton -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_rmsnorm.py b/test/registered/jit/test_rmsnorm.py index 824d77704c68..63ae38a16be1 100644 --- a/test/registered/jit/test_rmsnorm.py +++ b/test/registered/jit/test_rmsnorm.py @@ -4,7 +4,7 @@ import pytest import torch -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.srt.utils import is_hip from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/jit/test_rmsnorm_hf.py b/test/registered/jit/test_rmsnorm_hf.py index 8152e94568d6..c6b1beb80815 100644 --- a/test/registered/jit/test_rmsnorm_hf.py +++ b/test/registered/jit/test_rmsnorm_hf.py @@ -10,7 +10,7 @@ is_supported_rmsnorm_hf_hidden_size, rmsnorm_hf, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit 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_rope.py b/test/registered/jit/test_rope.py index e1ceb7287a99..08981ebf15ad 100644 --- a/test/registered/jit/test_rope.py +++ b/test/registered/jit/test_rope.py @@ -4,7 +4,7 @@ import torch import triton -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.srt.utils import is_hip from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/jit/test_set_mla_kv_buffer.py b/test/registered/jit/test_set_mla_kv_buffer.py index 4062355eff95..a841806b77a8 100644 --- a/test/registered/jit/test_set_mla_kv_buffer.py +++ b/test/registered/jit/test_set_mla_kv_buffer.py @@ -7,7 +7,7 @@ can_use_set_mla_kv_buffer, set_mla_kv_buffer, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import 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 c731da48c6dd..8693692c7125 100644 --- a/test/registered/jit/test_store_cache.py +++ b/test/registered/jit/test_store_cache.py @@ -5,7 +5,7 @@ import torch from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range 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_symm_mem_all_gather.py b/test/registered/jit/test_symm_mem_all_gather.py index fd3341a074a7..10098e2403f5 100644 --- a/test/registered/jit/test_symm_mem_all_gather.py +++ b/test/registered/jit/test_symm_mem_all_gather.py @@ -27,7 +27,7 @@ import sglang.srt.distributed.parallel_state as ps from sglang.jit_kernel.tests.utils import multigpu_pytest_main -from sglang.jit_kernel.utils import cache_once, get_ci_test_range +from sglang.kernels.jit import cache_once, get_ci_test_range from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( all_gather_inner, create_state, diff --git a/test/registered/jit/test_timestep_embedding.py b/test/registered/jit/test_timestep_embedding.py index bb92a8e7d629..caaf7d8ae1e5 100644 --- a/test/registered/jit/test_timestep_embedding.py +++ b/test/registered/jit/test_timestep_embedding.py @@ -13,7 +13,7 @@ from sglang.jit_kernel.timestep_embedding import ( timestep_embedding as timestep_embedding_cuda, ) -from sglang.jit_kernel.utils import get_ci_test_range +from sglang.kernels.jit import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=16, 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 6eccd8416dc8..6f6fd68acecf 100644 --- a/test/registered/jit/test_tp_qknorm.py +++ b/test/registered/jit/test_tp_qknorm.py @@ -21,7 +21,7 @@ ) from sglang.jit_kernel.mp import register_comm_cleanup from sglang.jit_kernel.tests.utils import multigpu_pytest_main -from sglang.jit_kernel.utils import cache_once +from sglang.kernels.jit import cache_once from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) diff --git a/test/registered/unit/distributed/test_vmm_utils.py b/test/registered/unit/distributed/test_vmm_utils.py index 8637969cfbc6..d56408442bd0 100644 --- a/test/registered/unit/distributed/test_vmm_utils.py +++ b/test/registered/unit/distributed/test_vmm_utils.py @@ -21,7 +21,7 @@ from cuda.bindings import driver as drv from sglang.jit_kernel.tests.utils import multigpu_pytest_main -from sglang.jit_kernel.utils import cache_once +from sglang.kernels.jit import cache_once from sglang.srt.distributed.device_communicators.vmm_utils import ( check_drv, exchange_posix_fds,