[Kernel] Phase 4 (layernorm): move JIT norm ops into kernels.ops.layernorm (RFC #29630)

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:

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,

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

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

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

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

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"])

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

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

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:

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:

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

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,

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:

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,

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:

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:

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,

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,

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,

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,

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

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,

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:

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,

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

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
 

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

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
 

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,

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,

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

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:

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

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

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

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:

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

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