[Kernel] Phase 3+4: move JIT infra + operator groups into sglang.kernels (RFC #29630) (#31666)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

test/registered/jit/benchmark/bench_custom_all_reduce.py

@@ -13,7 +13,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.mp import register_comm_cleanup
-from sglang.jit_kernel.utils import cache_once, is_arch_support_pdl
+from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py

@@ -13,7 +13,7 @@ from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import (
     dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm,
 )
 from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
-from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/benchmark/bench_dsv3_router_gemm.py

@@ -10,7 +10,7 @@ import torch.nn.functional as F
 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.utils import get_jit_cuda_arch, is_hip_runtime
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py

@@ -17,7 +17,7 @@ from sglang.jit_kernel.benchmark.utils import (
 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.utils import is_arch_support_pdl
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_set_mla_kv_buffer.py

@@ -21,7 +21,7 @@ from sglang.jit_kernel.benchmark.utils import (
     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.utils 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

test/registered/jit/benchmark/bench_symm_mem_all_gather.py

@@ -28,7 +28,7 @@ import torch.distributed as dist
 import sglang.srt.distributed.parallel_state as ps
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, multigpu_bench_main
-from sglang.jit_kernel.utils import cache_once
+from sglang.kernels.jit.utils import cache_once
 from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
     all_gather_inner,
     create_state,

test/registered/jit/benchmark/bench_tp_qknorm.py

@@ -32,7 +32,7 @@ from sglang.jit_kernel.all_reduce import (
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import multigpu_bench_main
 from sglang.jit_kernel.mp import register_comm_cleanup
-from sglang.jit_kernel.utils import cache_once, get_ci_test_range
+from sglang.kernels.jit.utils import cache_once, get_ci_test_range
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )

test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py

@@ -11,7 +11,7 @@ import flashinfer
 import torch
 
 from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE
-from sglang.jit_kernel.utils import KERNEL_PATH
+from sglang.kernels.jit.utils import KERNEL_PATH
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/benchmark/diffusion/bench_norm_impls.py

@@ -18,7 +18,7 @@ from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
 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.utils import KERNEL_PATH
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py

@@ -9,7 +9,7 @@ from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import (
 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.utils 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")

test/registered/jit/diffusion/test_qknorm_rope.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils 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")

test/registered/jit/diffusion/test_qwen_image_modulation.py

@@ -9,7 +9,7 @@ from sglang.jit_kernel.diffusion.triton.scale_shift import (
     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.utils 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")

test/registered/jit/diffusion/test_varlen_pack_pad.py

@@ -12,7 +12,7 @@ from sglang.jit_kernel.diffusion.triton.varlen_pack_pad import (
     fused_pack_qkv,
     fused_scatter_to_padded,
 )
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py

@@ -20,7 +20,7 @@ from sglang.jit_kernel.diffusion.triton.varlen_pack_pad import (
     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.utils import get_ci_test_range
 from sglang.multimodal_gen.runtime.layers.attention.backends import (
     flash_attn as _fa_backend,
 )

test/registered/jit/test_activation.py

@@ -9,7 +9,7 @@ from sglang.jit_kernel.activation import (
     relu2,
     run_activation,
 )
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_custom_all_reduce.py

@@ -31,7 +31,7 @@ import sglang.srt.distributed.parallel_state as ps
 from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_all_reduce_module
 from sglang.jit_kernel.mp import register_comm_cleanup
 from sglang.jit_kernel.tests.utils import multigpu_pytest_main
-from sglang.jit_kernel.utils import cache_once, get_ci_test_range
+from sglang.kernels.jit.utils import cache_once, get_ci_test_range
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )

test/registered/jit/test_cutedsl_bf16_gemm.py

@@ -5,7 +5,11 @@ import sys
 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.utils 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")

test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py

@@ -5,7 +5,11 @@ import sys
 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.utils 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")

test/registered/jit/test_dsv3_fused_a_gemm.py

@@ -7,7 +7,11 @@ import torch
 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.utils 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")

test/registered/jit/test_dsv3_router_gemm.py

@@ -7,7 +7,11 @@ import pytest
 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.utils 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")

test/registered/jit/test_fused_add_rmsnorm.py

@@ -4,7 +4,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils 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")

test/registered/jit/test_mla_kv_pack_quantize_fp8.py

@@ -4,7 +4,7 @@ import pytest
 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.utils 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")

test/registered/jit/test_moe_align_block_size.py

@@ -7,7 +7,7 @@ import triton
 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.utils 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")

test/registered/jit/test_moe_fused_gate.py

@@ -22,7 +22,7 @@ import pytest
 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.utils 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")

test/registered/jit/test_per_tensor_quant_fp8.py

@@ -6,7 +6,7 @@ import pytest
 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.utils 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")

test/registered/jit/test_per_token_group_quant.py

@@ -22,7 +22,7 @@ import pytest
 import torch
 
 from sglang.jit_kernel.per_token_group_quant import per_token_group_quant
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils import get_ci_test_range
 from sglang.kernels.ops.quantization.fp8_kernel import (
     create_per_token_group_quant_fp8_output_scale,
     fp8_dtype,

test/registered/jit/test_per_token_group_quant_8bit_v2.py

@@ -6,7 +6,7 @@ import torch
 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.utils 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")

test/registered/jit/test_qknorm.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils 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")

test/registered/jit/test_qknorm_across_heads.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils 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")

test/registered/jit/test_rmsnorm.py

@@ -4,7 +4,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils 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
 

test/registered/jit/test_rmsnorm_hf.py

@@ -10,7 +10,7 @@ from sglang.jit_kernel.rmsnorm_hf import (
     is_supported_rmsnorm_hf_hidden_size,
     rmsnorm_hf,
 )
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_rope.py

@@ -4,7 +4,7 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils 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
 

test/registered/jit/test_set_mla_kv_buffer.py

@@ -7,7 +7,7 @@ from sglang.jit_kernel.set_mla_kv_buffer import (
     can_use_set_mla_kv_buffer,
     set_mla_kv_buffer,
 )
-from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.kernels.jit.utils 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")

test/registered/jit/test_store_cache.py

@@ -5,7 +5,7 @@ import pytest
 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.utils 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")

test/registered/jit/test_symm_mem_all_gather.py

@@ -27,7 +27,7 @@ import torch.distributed as dist
 
 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.utils import cache_once, get_ci_test_range
 from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
     all_gather_inner,
     create_state,

test/registered/jit/test_timestep_embedding.py

@@ -13,7 +13,7 @@ except Exception:
 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.utils 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")

test/registered/jit/test_tp_qknorm.py

@@ -21,7 +21,7 @@ from sglang.jit_kernel.all_reduce import (
 )
 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.utils import cache_once
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )

test/registered/unit/distributed/test_vmm_utils.py

@@ -21,7 +21,7 @@ import torch.distributed as dist
 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.utils import cache_once
 from sglang.srt.distributed.device_communicators.vmm_utils import (
     check_drv,
     exchange_posix_fds,