[Kernel] Phase 4 batch-3: migrate tangled JIT subsystems + new groups into kernels.ops (RFC #29630) (#32045)

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

test/registered/attention/unittests/dual_chunk/test_dual_chunk_flash_attn.py

@@ -27,7 +27,7 @@ from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner i
 
 # Container gate (KNOWN_FAILURES.md §1): `DualChunkFlashAttentionBackend` calls
 # `flash_attn_varlen_func` on every forward via
-# `sglang.jit_kernel.flash_attention`. On SM8x/SM9x, that resolves to sgl-kernel's
+# `sglang.kernels.ops.attention.flash_attention`. On SM8x/SM9x, that resolves to sgl-kernel's
 # FA3 build (which works). On SM != {8, 9} (notably SM10.3 / GB300), the JIT
 # kernel falls back to the upstream `flash_attn` (FA2) wheel — but the
 # `lmsysorg/sglang:nightly-dev-cu13` container's `flash_attn` package ships
@@ -40,7 +40,7 @@ def _dual_chunk_fa_supported() -> tuple[bool, str]:
         return False, "CUDA is required"
     major, _minor = torch.cuda.get_device_capability()
     # FA3 path is taken when sm major is 8 or 9 (see
-    # `sglang.jit_kernel.flash_attention_v3._is_fa3_supported`). On that path
+    # `sglang.kernels.ops.attention.flash_attention_v3._is_fa3_supported`). On that path
     # the upstream `flash_attn` fallback is never invoked.
     if major in (8, 9):
         return True, ""

test/registered/eplb/test_lplb_distributed.py

@@ -82,7 +82,7 @@ def test_dispatch_probability_matches_torch_reference():
     """The fused CUDA `dispatch_probability` and the pure-torch reference
     must produce identical outputs for the same ``random_vals``. Single-rank,
     runs on any CUDA GPU."""
-    from sglang.jit_kernel.lplb.cuda_solver import (
+    from sglang.kernels.ops.lplb.cuda_solver import (
         dispatch_probability,
         dispatch_probability_torch_reference,
     )
@@ -135,8 +135,8 @@ def test_solve_ipm_matches_torch_reference():
     convergence test (the Big-M slack must reach ~0) and both backends would
     just return the 0.5 non-convergence sentinel — agreeing trivially without
     exercising the solve. Single-rank, any CUDA GPU."""
-    from sglang.jit_kernel.lplb.cuda_solver import solve_ipm as cuda_solve_ipm
-    from sglang.jit_kernel.lplb.torch_solver import solve_ipm_torch_reference
+    from sglang.kernels.ops.lplb.cuda_solver import solve_ipm as cuda_solve_ipm
+    from sglang.kernels.ops.lplb.torch_solver import solve_ipm_torch_reference
     from sglang.srt.eplb.lplb_solver import LPLBSolver
 
     torch.manual_seed(0)

test/registered/gemm/test_linear_bf16_fp32_hpc.py

@@ -1,6 +1,6 @@
 """Numerical tests for the HPC-Ops bf16xfp32 router GEMM path.
 
-Validates sglang.jit_kernel.dsv4.linear_bf16_fp32's HPC-Ops branch against
+Validates sglang.kernels.ops.attention.dsv4.linear_bf16_fp32's HPC-Ops branch against
 the fp32 reference on the LongCat-Flash router shapes. Skipped when HPC-Ops
 (https://github.com/Tencent/hpc-ops) is not installed or the GPU is not
 Hopper (the kernels ship sm90a only).
@@ -10,7 +10,7 @@ import unittest
 
 import torch
 
-from sglang.jit_kernel.dsv4.gemm import (
+from sglang.kernels.ops.attention.dsv4.gemm import (
     _hpc_gemm_bf16xfp32_available,
     _linear_bf16_fp32_hpc,
     linear_bf16_fp32,

test/registered/jit/benchmark/bench_add_constant.py

@@ -2,12 +2,15 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.add_constant import _jit_add_constant_module, add_constant
 from sglang.jit_kernel.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark_no_cudagraph,
 )
+from sglang.kernels.ops.attention.add_constant import (
+    _jit_add_constant_module,
+    add_constant,
+)
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_clamp_position.py

@@ -9,7 +9,7 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
-from sglang.jit_kernel.clamp_position import clamp_position_cuda
+from sglang.kernels.ops.attention.clamp_position import clamp_position_cuda
 from sglang.srt.utils import get_compiler_backend
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/benchmark/bench_concat_mla.py

@@ -7,8 +7,8 @@ from sgl_kernel import concat_mla_absorb_q as aot_absorb_q
 from sgl_kernel import concat_mla_k as aot_k
 
 from sglang.jit_kernel.benchmark.utils import run_benchmark
-from sglang.jit_kernel.concat_mla import concat_mla_absorb_q as jit_absorb_q
-from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k
+from sglang.kernels.ops.attention.concat_mla import concat_mla_absorb_q as jit_absorb_q
+from sglang.kernels.ops.attention.concat_mla import concat_mla_k as jit_k
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/benchmark/bench_fused_qknorm_rope.py

@@ -15,7 +15,7 @@ import triton
 import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
-from sglang.jit_kernel.fused_qknorm_rope import (
+from sglang.kernels.ops.attention.fused_qknorm_rope import (
     fused_qk_norm_rope as fused_qk_norm_rope_jit,
 )
 from sglang.test.ci.ci_register import register_cuda_ci

test/registered/jit/benchmark/bench_hadamard.py

@@ -13,7 +13,7 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
-from sglang.jit_kernel.hadamard import hadamard_transform
+from sglang.kernels.ops.attention.hadamard import hadamard_transform
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py

@@ -14,10 +14,10 @@ from sglang.jit_kernel.benchmark.utils import (
     DEFAULT_QUANTILES,
     get_benchmark_range,
 )
-from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import (
+from sglang.kernels.jit.utils import is_arch_support_pdl
+from sglang.kernels.ops.attention.mla_kv_pack_quantize_fp8 import (
     mla_kv_pack_quantize_fp8 as hybrid_pack,
 )
-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_moe_fused_gate.py

@@ -2,7 +2,7 @@ import torch
 
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import create_random
-from sglang.jit_kernel.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit
+from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_online_c128_mtp.py

@@ -15,7 +15,9 @@ from sglang.jit_kernel.benchmark.utils import (
     run_benchmark,
     run_benchmark_no_cudagraph,
 )
-from sglang.jit_kernel.dsv4.online_c128_mtp import _jit_online_c128_mtp_module
+from sglang.kernels.ops.attention.dsv4.online_c128_mtp import (
+    _jit_online_c128_mtp_module,
+)
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_rope.py

@@ -75,7 +75,7 @@ def sglang_pos_enc_rope(
     positions: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.jit_kernel.rope import rotary_embedding_with_key
+    from sglang.kernels.ops.attention.rope import rotary_embedding_with_key
 
     head_size = q.shape[-1]
     rotary_embedding_with_key(
@@ -94,7 +94,7 @@ def sglang_fused_rope(
     positions: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.jit_kernel.rope import apply_rope_inplace
+    from sglang.kernels.ops.attention.rope import apply_rope_inplace
 
     apply_rope_inplace(q, k, COS_SIN_CACHE, positions, is_neox=is_neox)
 
@@ -114,7 +114,7 @@ def jit_rope_then_store(
     out_loc: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.jit_kernel.rope import apply_rope_inplace
+    from sglang.kernels.ops.attention.rope import apply_rope_inplace
     from sglang.kernels.ops.kvcache.kvcache import store_cache
 
     head_size = q.shape[-1]
@@ -146,7 +146,7 @@ def jit_fused_rope_store(
     out_loc: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache
+    from sglang.kernels.ops.attention.rope import apply_rope_inplace_with_kvcache
 
     apply_rope_inplace_with_kvcache(
         q, k, v, k_cache, v_cache, COS_SIN_CACHE, positions, out_loc, is_neox=is_neox

test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py

@@ -7,7 +7,9 @@ import triton
 import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
-from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import sparse_mla_q8kv8_prefill_fwd
+from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
+    sparse_mla_q8kv8_prefill_fwd,
+)
 from sglang.srt.utils import is_sm90_supported
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

test/registered/jit/benchmark/bench_topk.py

@@ -1,7 +1,7 @@
 import torch
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.dsv4.topk import (
+from sglang.kernels.ops.attention.dsv4.topk import (
     plan_topk_v2,
     topk_transform_512,
     topk_transform_512_v2,

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

@@ -3,10 +3,10 @@ from dataclasses import dataclass
 import torch
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import (
+from sglang.kernels.ops.diffusion.causal_conv3d_cat_pad import (
     fused_causal_conv3d_cat_pad_cuda,
 )
-from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import (
+from sglang.kernels.ops.diffusion.triton.causal_conv3d_pad import (
     fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

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

@@ -10,7 +10,7 @@ import torch
 import torch.nn.functional as F
 import triton.testing
 
-from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu
+from sglang.kernels.ops.diffusion.triton.group_norm_silu import triton_group_norm_silu
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 

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

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import torch
 
-from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
+from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
     ltx2_qknorm_split_rope_cuda,
 )
 from sglang.test.ci.ci_register import register_cuda_ci

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

@@ -14,11 +14,11 @@ import torch
 import torch.nn.functional as F
 
 from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE
-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.kernels.jit.utils import KERNEL_PATH
+from sglang.kernels.ops.diffusion.triton.norm import norm_infer, rms_norm_fn
+from sglang.kernels.ops.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 

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

@@ -153,7 +153,7 @@ def split_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
 
 
 def fused_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
-    from sglang.jit_kernel.diffusion.qknorm_rope import fused_inplace_qknorm_rope
+    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
 
     fused_inplace_qknorm_rope(
         inputs["q"],

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

@@ -4,8 +4,8 @@ import torch
 import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
-from sglang.jit_kernel.diffusion.triton.norm import norm_infer
-from sglang.jit_kernel.diffusion.triton.scale_shift import (
+from sglang.kernels.ops.diffusion.triton.norm import norm_infer
+from sglang.kernels.ops.diffusion.triton.scale_shift import (
     fuse_layernorm_scale_shift_gate_select01_kernel,
     fuse_residual_layernorm_scale_shift_gate_select01_kernel,
 )

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

@@ -4,8 +4,8 @@ from dataclasses import dataclass
 
 import torch
 
-from sglang.jit_kernel.diffusion.residual_gate_add import residual_gate_add_cuda
-from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
+from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add_cuda
+from sglang.kernels.ops.diffusion.triton.scale_shift import fuse_scale_shift_kernel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/benchmark/kv_canary/bench_plan.py

@@ -20,9 +20,9 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
-from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
-from sglang.jit_kernel.kv_canary.verify import VerifyPlan
-from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels
+from sglang.kernels.ops.kv_canary.verify import VerifyPlan
+from sglang.kernels.ops.kv_canary.write import WritePlan
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True)

test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py

@@ -11,7 +11,7 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark_no_cudagraph,
 )
-from sglang.jit_kernel.kv_canary.scatter_req_token_ids import (
+from sglang.kernels.ops.kv_canary.scatter_req_token_ids import (
     launch_scatter_req_token_ids_kernel,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

test/registered/jit/benchmark/kv_canary/bench_verify.py

@@ -21,8 +21,8 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.verify import (
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.verify import (
     CANARY_SLOT_BYTES,
     CanaryLaunchTag,
     RealKvSource,

test/registered/jit/benchmark/kv_canary/bench_write.py

@@ -21,13 +21,13 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.verify import (
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.verify import (
     CANARY_SLOT_BYTES,
     CanaryLaunchTag,
     VerifyOrWriteContext,
 )
-from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
+from sglang.kernels.ops.kv_canary.write import WritePlan, launch_canary_write_kernel
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True)

test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py

@@ -10,7 +10,7 @@ import torch
 import triton
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk
+from sglang.kernels.ops.attention.minimax_decode_topk import minimax_decode_topk
 from sglang.kernels.ops.attention.minimax_sparse.decode.flash_with_topk_idx import (
     _topk_index_merge_kernel,
     _topk_index_partial_kernel,

test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py

@@ -6,7 +6,7 @@ launch) vs the unfused path (GemmaRMSNorm(q) + GemmaRMSNorm(k) + rotary_emb,
 import torch
 
 from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.minimax_qknorm_rope import (
+from sglang.kernels.ops.attention.minimax_qknorm_rope import (
     minimax_qknorm_rope,
     minimax_qknorm_rope_grouped,
 )

test/registered/jit/deepseek_v4/test_c128_v2.py

@@ -7,7 +7,6 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.dsv4 import compress_forward
 from sglang.jit_kernel.tests.deepseek_v4.common import (
     LegacyContext,
     PagedContext,
@@ -16,6 +15,7 @@ from sglang.jit_kernel.tests.deepseek_v4.common import (
     make_state_pool,
     to_seq_extend,
 )
+from sglang.kernels.ops.attention.dsv4 import compress_forward
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/deepseek_v4/test_c4_v2.py

@@ -7,7 +7,6 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.dsv4 import compress_forward
 from sglang.jit_kernel.tests.deepseek_v4.common import (
     LegacyContext,
     PagedContext,
@@ -16,6 +15,7 @@ from sglang.jit_kernel.tests.deepseek_v4.common import (
     make_state_pool,
     to_seq_extend,
 )
+from sglang.kernels.ops.attention.dsv4 import compress_forward
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/deepseek_v4/test_fp4_indexer.py

@@ -5,15 +5,15 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.dsv4 import (
-    CompressorDecodePlan,
-    compress_norm_rope_store,
-    fused_q_indexer_rope_hadamard_fp4_quant,
-)
 from sglang.kernels.ops.attention.deepseek_v4_rope import (
     apply_rotary_emb_triton,
     precompute_freqs_cis,
 )
+from sglang.kernels.ops.attention.dsv4 import (
+    CompressorDecodePlan,
+    compress_norm_rope_store,
+    fused_q_indexer_rope_hadamard_fp4_quant,
+)
 from sglang.kernels.ops.attention.dsv4.fp4_indexer import (
     quantize_fp4_indexer_tensor,
     store_fp4_index_k_cache,
@@ -27,7 +27,7 @@ _is_xpu = is_xpu()
 if _is_xpu:
     from sgl_kernel import hadamard_transform
 else:
-    from sglang.jit_kernel.hadamard import hadamard_transform
+    from sglang.kernels.ops.attention.hadamard import hadamard_transform
 
 HEAD_DIM = 128
 FP4_DIM = HEAD_DIM // 2

test/registered/jit/deepseek_v4/test_fp8_wo_a.py

@@ -9,8 +9,8 @@ import unittest
 
 import torch
 
-import sglang.jit_kernel.dsv4.fp8_wo_a as fp8_wo_a_module
-from sglang.jit_kernel.dsv4 import sglang_per_token_group_quant_fp8_dsv4_wo_a
+import sglang.kernels.ops.attention.dsv4.fp8_wo_a as fp8_wo_a_module
+from sglang.kernels.ops.attention.dsv4 import sglang_per_token_group_quant_fp8_dsv4_wo_a
 from sglang.kernels.ops.quantization.fp8_kernel import (
     fp8_dtype,
     sglang_per_token_group_quant_fp8,

test/registered/jit/deepseek_v4/test_topk_v2.py

@@ -29,7 +29,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.dsv4.topk import plan_topk_v2, topk_transform_512_v2
+from sglang.kernels.ops.attention.dsv4.topk import plan_topk_v2, topk_transform_512_v2
 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/diffusion/test_causal_conv3d_cat_pad.py

@@ -3,13 +3,13 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.diffusion.causal_conv3d_cat_pad import (
     fused_causal_conv3d_cat_pad_cuda,
 )
-from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import (
+from sglang.kernels.ops.diffusion.triton.causal_conv3d_pad import (
     fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
 )
-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_flydsl_fused_norm.py

@@ -58,7 +58,7 @@ FUSED_CASES = [
 
 @pytest.mark.parametrize("norm_type,B,L", FUSED_CASES)
 def test_fused_residual_norm_scale_shift(norm_type, B, L):
-    from sglang.jit_kernel.diffusion.flydsl.fused_residual_norm import (
+    from sglang.kernels.ops.diffusion.flydsl.fused_residual_norm import (
         flydsl_fused_residual_norm_scale_shift,
     )
 
@@ -110,7 +110,7 @@ NSS_CASES = [
 
 @pytest.mark.parametrize("norm_type,B,L", NSS_CASES)
 def test_norm_scale_shift(norm_type, B, L):
-    from sglang.jit_kernel.diffusion.flydsl.fused_residual_norm import (
+    from sglang.kernels.ops.diffusion.flydsl.fused_residual_norm import (
         flydsl_norm_scale_shift,
     )
 

test/registered/jit/diffusion/test_fused_norm_scale_shift.py

@@ -6,7 +6,7 @@ import torch
 from einops import rearrange
 from torch import Tensor
 
-from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import (
+from sglang.kernels.ops.diffusion.cutedsl.scale_residual_norm_scale_shift import (
     fused_norm_scale_shift,
     fused_scale_residual_norm_scale_shift,
     validate_scale_shift,

test/registered/jit/diffusion/test_group_norm_silu.py

@@ -5,8 +5,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu
-from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu
+from sglang.kernels.ops.diffusion.group_norm_silu import apply_group_norm_silu
+from sglang.kernels.ops.diffusion.triton.group_norm_silu import triton_group_norm_silu
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/diffusion/test_ltx2_ada_values.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.diffusion.triton.ltx2_ada_values import ltx2_ada_values9
+from sglang.kernels.ops.diffusion.triton.ltx2_ada_values import ltx2_ada_values9
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py

@@ -4,7 +4,7 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
+from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
     can_use_ltx2_qknorm_split_rope_cuda,
     ltx2_qknorm_split_rope_cuda,
 )

test/registered/jit/diffusion/test_qknorm_rope.py

@@ -70,7 +70,7 @@ def fused_qknorm_rope(
     positions: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.jit_kernel.diffusion.qknorm_rope import fused_inplace_qknorm_rope
+    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
 
     fused_inplace_qknorm_rope(
         q,

test/registered/jit/diffusion/test_qwen_image_modulation.py

@@ -4,12 +4,12 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.diffusion.triton.norm import norm_infer
-from sglang.jit_kernel.diffusion.triton.scale_shift import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.diffusion.triton.norm import norm_infer
+from sglang.kernels.ops.diffusion.triton.scale_shift import (
     fuse_layernorm_scale_shift_gate_select01_kernel,
     fuse_residual_layernorm_scale_shift_gate_select01_kernel,
 )
-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_residual_gate_add.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.diffusion.residual_gate_add import (
+from sglang.kernels.ops.diffusion.residual_gate_add import (
     can_use_residual_gate_add_cuda,
     residual_gate_add_cuda,
 )

test/registered/jit/diffusion/test_varlen_pack_pad.py

@@ -7,12 +7,12 @@ zeros + index_copy_) across bf16/fp16 and several shape/mask cases.
 import pytest
 import torch
 
-from sglang.jit_kernel.diffusion.triton.varlen_pack_pad import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.diffusion.triton.varlen_pack_pad import (
     build_inv_indices,
     fused_pack_qkv,
     fused_scatter_to_padded,
 )
-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

@@ -15,12 +15,12 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from sglang.jit_kernel.diffusion.triton.varlen_pack_pad import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
+from sglang.kernels.ops.diffusion.triton.varlen_pack_pad import (
     fused_pack_qkv,
     fused_scatter_to_padded,
 )
-from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
-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/kv_canary/test_const_sync.py

@@ -4,7 +4,7 @@ import re
 from pathlib import Path
 
 import sglang.jit_kernel
-from sglang.jit_kernel.kv_canary import consts
+from sglang.kernels.ops.kv_canary import consts
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/kv_canary/test_kernel_config.py

@@ -3,14 +3,6 @@ from __future__ import annotations
 import pytest
 import torch
 
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.verify import (
-    CanaryLaunchTag,
-    VerifyOrWriteContext,
-    VerifyPlan,
-    launch_canary_verify_kernel,
-)
-from sglang.jit_kernel.kv_canary.write import WritePlan
 from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_canary_buf_equal,
@@ -33,6 +25,14 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
     dummy_pseudo_tensors,
     empty_extras,
 )
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.verify import (
+    CanaryLaunchTag,
+    VerifyOrWriteContext,
+    VerifyPlan,
+    launch_canary_verify_kernel,
+)
+from sglang.kernels.ops.kv_canary.write import WritePlan
 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/kv_canary/test_pipeline_e2e.py

@@ -5,25 +5,6 @@ from typing import Any, Optional
 import pytest
 import torch
 
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
-from sglang.jit_kernel.kv_canary.plan_ref import (
-    launch_canary_plan_kernels_torch_reference,
-)
-from sglang.jit_kernel.kv_canary.verify import (
-    CanaryLaunchTag,
-    RealKvSource,
-    VerifyOrWriteContext,
-    VerifyPlan,
-    launch_canary_verify_kernel,
-)
-from sglang.jit_kernel.kv_canary.verify_ref import (
-    launch_canary_verify_kernel_torch_reference,
-)
-from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
-from sglang.jit_kernel.kv_canary.write_ref import (
-    launch_canary_write_kernel_torch_reference,
-)
 from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_canary_buf_equal,
@@ -38,6 +19,25 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
     empty_extras,
     make_req_to_token,
 )
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels
+from sglang.kernels.ops.kv_canary.plan_ref import (
+    launch_canary_plan_kernels_torch_reference,
+)
+from sglang.kernels.ops.kv_canary.verify import (
+    CanaryLaunchTag,
+    RealKvSource,
+    VerifyOrWriteContext,
+    VerifyPlan,
+    launch_canary_verify_kernel,
+)
+from sglang.kernels.ops.kv_canary.verify_ref import (
+    launch_canary_verify_kernel_torch_reference,
+)
+from sglang.kernels.ops.kv_canary.write import WritePlan, launch_canary_write_kernel
+from sglang.kernels.ops.kv_canary.write_ref import (
+    launch_canary_write_kernel_torch_reference,
+)
 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/kv_canary/test_plan_hand.py

@@ -5,12 +5,6 @@ import random
 import pytest
 import torch
 
-from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
-from sglang.jit_kernel.kv_canary.plan_ref import (
-    launch_canary_plan_kernels_torch_reference,
-)
-from sglang.jit_kernel.kv_canary.verify import VerifyPlan
-from sglang.jit_kernel.kv_canary.write import WritePlan
 from sglang.jit_kernel.tests.kv_canary._differential import run_plan_diff
 from sglang.jit_kernel.tests.kv_canary._fixtures import (
     allocate_plan_pair,
@@ -20,6 +14,12 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
     make_req_to_token,
 )
 from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
+from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels
+from sglang.kernels.ops.kv_canary.plan_ref import (
+    launch_canary_plan_kernels_torch_reference,
+)
+from sglang.kernels.ops.kv_canary.verify import VerifyPlan
+from sglang.kernels.ops.kv_canary.write import WritePlan
 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/kv_canary/test_scatter_req_token_ids.py

@@ -5,7 +5,7 @@ import unittest
 
 import torch
 
-from sglang.jit_kernel.kv_canary.scatter_req_token_ids import (
+from sglang.kernels.ops.kv_canary.scatter_req_token_ids import (
     _SCATTER_BATCH_BLOCK,
     launch_scatter_req_token_ids_kernel,
     scatter_req_token_ids_torch_reference,

test/registered/jit/kv_canary/test_verify_fuzz.py

@@ -6,12 +6,6 @@ from dataclasses import dataclass
 import pytest
 import torch
 
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.verify import (
-    CanaryLaunchTag,
-    RealKvSource,
-    VerifyPlan,
-)
 from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     make_canary_buf,
@@ -29,6 +23,12 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
     run_fuzz_combo,
 )
 from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.verify import (
+    CanaryLaunchTag,
+    RealKvSource,
+    VerifyPlan,
+)
 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/kv_canary/test_verify_hand.py

@@ -8,22 +8,6 @@ from typing import Callable
 import pytest
 import torch
 
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
-from sglang.jit_kernel.kv_canary.verify import (
-    CanaryLaunchTag,
-    RealKvSource,
-    VerifyOrWriteContext,
-    VerifyPlan,
-    launch_canary_verify_kernel,
-)
-from sglang.jit_kernel.kv_canary.verify_ref import (
-    _compute_real_kv_hash_scalar,
-    launch_canary_verify_kernel_torch_reference,
-)
-from sglang.jit_kernel.kv_canary.write_ref import (
-    launch_canary_write_kernel_torch_reference,
-)
 from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_only_bits_set,
@@ -51,6 +35,22 @@ from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
     _hand_fold_all,
     _hand_fold_partial,
 )
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3
+from sglang.kernels.ops.kv_canary.verify import (
+    CanaryLaunchTag,
+    RealKvSource,
+    VerifyOrWriteContext,
+    VerifyPlan,
+    launch_canary_verify_kernel,
+)
+from sglang.kernels.ops.kv_canary.verify_ref import (
+    _compute_real_kv_hash_scalar,
+    launch_canary_verify_kernel_torch_reference,
+)
+from sglang.kernels.ops.kv_canary.write_ref import (
+    launch_canary_write_kernel_torch_reference,
+)
 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/kv_canary/test_write_fuzz.py

@@ -6,12 +6,6 @@ from dataclasses import dataclass
 import pytest
 import torch
 
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.verify import (
-    CanaryLaunchTag,
-    RealKvSource,
-)
-from sglang.jit_kernel.kv_canary.write import WritePlan
 from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     make_canary_buf,
@@ -29,6 +23,12 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
     run_fuzz_combo,
 )
 from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.verify import (
+    CanaryLaunchTag,
+    RealKvSource,
+)
+from sglang.kernels.ops.kv_canary.write import WritePlan
 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/kv_canary/test_write_hand.py

@@ -7,19 +7,6 @@ from unittest.mock import patch
 import pytest
 import torch
 
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary import write as write_module
-from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
-from sglang.jit_kernel.kv_canary.verify import (
-    CANARY_SLOT_BYTES,
-    CanaryLaunchTag,
-    RealKvSource,
-    VerifyOrWriteContext,
-    launch_canary_verify_kernel,
-)
-from sglang.jit_kernel.kv_canary.write import (
-    launch_canary_write_kernel,
-)
 from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_canary_state_equal,
@@ -49,6 +36,19 @@ from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
     _hand_fold_all,
     _hand_fold_partial,
 )
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary import write as write_module
+from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3
+from sglang.kernels.ops.kv_canary.verify import (
+    CANARY_SLOT_BYTES,
+    CanaryLaunchTag,
+    RealKvSource,
+    VerifyOrWriteContext,
+    launch_canary_verify_kernel,
+)
+from sglang.kernels.ops.kv_canary.write import (
+    launch_canary_write_kernel,
+)
 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/minimax/test_minimax_decode_topk.py

@@ -9,7 +9,7 @@ the consumer ``_gqa_share_sparse_decode_kernel`` contract.
 import pytest
 import torch
 
-from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk
+from sglang.kernels.ops.attention.minimax_decode_topk import minimax_decode_topk
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/minimax/test_minimax_decode_topk_page_table.py

@@ -16,7 +16,7 @@ import torch
 
 flashinfer = pytest.importorskip("flashinfer")
 
-from sglang.jit_kernel.minimax_decode_topk import (
+from sglang.kernels.ops.attention.minimax_decode_topk import (
     minimax_decode_topk,
     minimax_decode_topk_page_table,
 )

test/registered/jit/minimax/test_minimax_qknorm_rope.py

@@ -8,7 +8,7 @@ convention (cos|sin cache, neox pairs (i, i+rotary_dim/2)).
 import pytest
 import torch
 
-from sglang.jit_kernel.minimax_qknorm_rope import (
+from sglang.kernels.ops.attention.minimax_qknorm_rope import (
     minimax_qknorm_rope,
     minimax_qknorm_rope_grouped,
 )

test/registered/jit/test_add_constant.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.add_constant import add_constant
+from sglang.kernels.ops.attention.add_constant import add_constant
 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/test_clamp_position.py

@@ -3,7 +3,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.clamp_position import clamp_position_cuda
+from sglang.kernels.ops.attention.clamp_position import clamp_position_cuda
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_concat_mla.py

@@ -58,7 +58,7 @@ def jit_concat_mla_k(
     k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor
 ) -> None:
     """JIT compiled implementation."""
-    from sglang.jit_kernel.concat_mla import concat_mla_k
+    from sglang.kernels.ops.attention.concat_mla import concat_mla_k
 
     concat_mla_k(k, k_nope, k_rope)
 
@@ -67,7 +67,7 @@ def jit_concat_mla_absorb_q(
     a: torch.Tensor, b: torch.Tensor, out: torch.Tensor
 ) -> None:
     """JIT compiled implementation - wrapper for test compatibility."""
-    from sglang.jit_kernel.concat_mla import concat_mla_absorb_q
+    from sglang.kernels.ops.attention.concat_mla import concat_mla_absorb_q
 
     result = concat_mla_absorb_q(a, b)
     out.copy_(result)

test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py

@@ -89,15 +89,15 @@ from typing import Callable, Iterable, Literal, Optional
 import pytest
 import torch
 
-from sglang.jit_kernel.dsv4 import (
-    CompressorDecodePlan,
-    CompressorPrefillPlan,
-    compress_forward,
-)
 from sglang.jit_kernel.tests.deepseek_v4.common import (
     make_legacy_context,
     to_seq_extend,
 )
+from sglang.kernels.ops.attention.dsv4 import (
+    CompressorDecodePlan,
+    CompressorPrefillPlan,
+    compress_forward,
+)
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/test_dsv32_indexer_fusion.py

@@ -17,15 +17,15 @@ from __future__ import annotations
 import pytest
 import torch
 
-from sglang.jit_kernel.dsv4 import fused_q_indexer_rope_first_quant
-from sglang.jit_kernel.dsv32 import (
-    fused_k_indexer_norm_rope,
-    fused_k_indexer_norm_rope_store,
-)
-from sglang.jit_kernel.fused_store_index_cache import (
+from sglang.kernels.ops.attention.dsv4 import fused_q_indexer_rope_first_quant
+from sglang.kernels.ops.attention.fused_store_index_cache import (
     can_use_dsa_fused_store,
     fused_store_index_k_cache,
 )
+from sglang.kernels.ops.quantization.dsv32 import (
+    fused_k_indexer_norm_rope,
+    fused_k_indexer_norm_rope_store,
+)
 from sglang.srt.utils import is_hip
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/test_flash_attention_3_only_qv.py

@@ -11,7 +11,7 @@ import pytest
 import torch
 from einops import rearrange, repeat
 
-from sglang.jit_kernel.flash_attention import flash_attn_with_kvcache
+from sglang.kernels.ops.attention.flash_attention import flash_attn_with_kvcache
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 

test/registered/jit/test_flash_attention_4.py

@@ -11,7 +11,7 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange, repeat
 
-from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
+from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_fused_metadata_copy.py

@@ -217,7 +217,9 @@ def test_fused_metadata_copy_dtype_validation():
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")
 
-    from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda
+    from sglang.kernels.ops.attention.fused_metadata_copy import (
+        fused_metadata_copy_cuda,
+    )
 
     bs = 2
     max_len = 128
@@ -330,7 +332,9 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")
 
-    from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda
+    from sglang.kernels.ops.attention.fused_metadata_copy import (
+        fused_metadata_copy_cuda,
+    )
 
     max_len = 128
     max_seqlen_k = 256
@@ -431,7 +435,9 @@ def test_fused_metadata_copy_large_batch(bs):
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")
 
-    from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda
+    from sglang.kernels.ops.attention.fused_metadata_copy import (
+        fused_metadata_copy_cuda,
+    )
 
     forward_mode = 0  # DECODE
     max_len = 128
@@ -626,7 +632,9 @@ def test_fused_metadata_copy_multi_dtype_validation():
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")
 
-    from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda
+    from sglang.kernels.ops.attention.fused_metadata_copy import (
+        fused_metadata_copy_multi_cuda,
+    )
 
     bs = 2
     max_len = 128
@@ -721,7 +729,9 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")
 
-    from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda
+    from sglang.kernels.ops.attention.fused_metadata_copy import (
+        fused_metadata_copy_multi_cuda,
+    )
 
     max_len = 128
     seqlens_expanded_size = bs
@@ -923,7 +933,9 @@ def test_fused_metadata_copy_multi_large_batch(bs):
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")
 
-    from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda
+    from sglang.kernels.ops.attention.fused_metadata_copy import (
+        fused_metadata_copy_multi_cuda,
+    )
 
     max_len = 128
     seqlens_expanded_size = bs

test/registered/jit/test_fused_store_index_cache.py

@@ -25,7 +25,7 @@ import torch
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 try:
-    from sglang.jit_kernel.fused_store_index_cache import (
+    from sglang.kernels.ops.attention.fused_store_index_cache import (
         can_use_dsa_fused_store,
         fused_store_index_k_cache,
     )

test/registered/jit/test_hadamard_jit.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn.functional as F
 from scipy.linalg import hadamard
 
-from sglang.jit_kernel.hadamard import (
+from sglang.kernels.ops.attention.hadamard import (
     hadamard_transform,
     hadamard_transform_12n,
     hadamard_transform_20n,

test/registered/jit/test_inkling_attn_prologue_tau.py

@@ -7,7 +7,9 @@ must be untouched by tau.
 import pytest
 import torch
 
-from sglang.jit_kernel.inkling_attn_prologue import inkling_attn_prologue_decode
+from sglang.kernels.ops.model.inkling.inkling_attn_prologue import (
+    inkling_attn_prologue_decode,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_inkling_rel_proj.py

@@ -6,7 +6,7 @@ production strided-r layout and contiguous inputs."""
 import pytest
 import torch
 
-from sglang.jit_kernel.inkling_rel_proj import rel_proj_small_t
+from sglang.kernels.ops.model.inkling.inkling_rel_proj import rel_proj_small_t
 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_inkling_row_scale.py

@@ -5,10 +5,10 @@ including on the row-strided qkvr-slice layouts."""
 import pytest
 import torch
 
-from sglang.jit_kernel.inkling_row_scale import row_scale_bf16
 from sglang.kernels.ops.attention.log_scaling_tau import (
     _apply_log_scaling_tau_kernel,
 )
+from sglang.kernels.ops.model.inkling.inkling_row_scale import row_scale_bf16
 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")
@@ -51,7 +51,7 @@ def test_row_compact_bitexact(rows, inner, strided):
     """The tau-less compaction flavor (kHasTau=false) must reproduce
     .contiguous() exactly on the same strided layouts row_scale handles --
     no other test exercises run_compact."""
-    from sglang.jit_kernel.inkling_row_scale import row_compact_bf16
+    from sglang.kernels.ops.model.inkling.inkling_row_scale import row_compact_bf16
 
     torch.manual_seed(rows + inner)
     if strided:

test/registered/jit/test_kpool_topk_transform.py

@@ -14,7 +14,7 @@ from typing import Optional
 import pytest
 import torch
 
-from sglang.jit_kernel.kpool_topk_transform import fast_kpool_topk_transform_fused
+from sglang.kernels.ops.moe.kpool_topk_transform import fast_kpool_topk_transform_fused
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_minimax_m3_qk_norm_rope.py

@@ -14,7 +14,7 @@ if not is_hip():
 if not torch.cuda.is_available():
     pytest.skip("Requires a GPU.", allow_module_level=True)
 
-from sglang.jit_kernel.minimax_m3.qk_norm_rope import (  # noqa: E402
+from sglang.kernels.ops.attention.minimax_m3_qk_norm_rope import (  # noqa: E402
     qk_gemma_rmsnorm_rope,
     sparse_qk_index_gemma_rmsnorm_rope,
     sparse_qk_index_gemma_rmsnorm_rope_cache,

test/registered/jit/test_mla_kv_pack_quantize_fp8.py

@@ -3,8 +3,10 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import mla_kv_pack_quantize_fp8
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.attention.mla_kv_pack_quantize_fp8 import (
+    mla_kv_pack_quantize_fp8,
+)
 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

@@ -6,8 +6,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.jit_kernel.moe_align import moe_align_block_size
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.moe.moe_align import moe_align_block_size
 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

@@ -21,8 +21,8 @@ from typing import Tuple
 import pytest
 import torch
 
-from sglang.jit_kernel.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit
 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_moe_lora_align_block_size.py

@@ -8,7 +8,7 @@ import torch
 # ---------------------------------------------------------
 # IMPORT PREBUILT KERNEL
 # ---------------------------------------------------------
-from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
+from sglang.kernels.ops.moe.moe_lora_align import moe_lora_align_block_size
 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_wna16_marlin.py

@@ -6,7 +6,7 @@ import pytest
 import torch
 from sgl_kernel.scalar_type import scalar_types
 
-from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
+from sglang.kernels.ops.moe.moe_wna16_marlin import moe_wna16_marlin_gemm
 from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
 from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import fused_marlin_moe
 from sglang.srt.layers.quantization.marlin_utils_fp4 import (

test/registered/jit/test_pos_enc.py

@@ -7,7 +7,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.jit_kernel.rope import rotary_embedding
+from sglang.kernels.ops.attention.rope import rotary_embedding
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=18, stage="base-b-kernel-unit", runner_config="1-gpu-large")

test/registered/jit/test_rope.py

@@ -53,7 +53,7 @@ def sglang_jit_rope(
     positions: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.jit_kernel.rope import apply_rope_inplace
+    from sglang.kernels.ops.attention.rope import apply_rope_inplace
 
     apply_rope_inplace(q, k, cos_sin_cache, positions, is_neox=is_neox)
 
@@ -250,7 +250,7 @@ def test_fused_rope_store(
     is_neox: bool,
 ) -> None:
     """Test fused RoPE + KV cache store against separate RoPE + manual store."""
-    from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache
+    from sglang.kernels.ops.attention.rope import apply_rope_inplace_with_kvcache
 
     num_qo_heads = num_kv_heads * gqa_ratio
     dtype = DTYPE

test/registered/jit/test_sigmoid_gate_mul.py

@@ -28,7 +28,7 @@ def reference_sigmoid_gate_mul(x, gate):
     ],
 )
 def test_sigmoid_gate_mul_correctness(shape, dtype):
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import sigmoid_gate_mul
 
     torch.manual_seed(42)
     x = torch.randn(shape, dtype=dtype, device=DEVICE)
@@ -44,7 +44,7 @@ def test_sigmoid_gate_mul_correctness(shape, dtype):
 
 @pytest.mark.parametrize("shape", [(4, 4096), (1, 128)])
 def test_sigmoid_gate_mul_does_not_modify_inputs(shape):
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import sigmoid_gate_mul
 
     torch.manual_seed(42)
     x = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE)
@@ -59,7 +59,7 @@ def test_sigmoid_gate_mul_does_not_modify_inputs(shape):
 
 
 def test_sigmoid_gate_mul_output_dtype():
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import sigmoid_gate_mul
 
     for dtype in [torch.bfloat16, torch.float16, torch.float32]:
         x = torch.randn(4, 4096, dtype=dtype, device=DEVICE)
@@ -69,7 +69,7 @@ def test_sigmoid_gate_mul_output_dtype():
 
 
 def test_sigmoid_gate_mul_contiguous_output():
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import sigmoid_gate_mul
 
     x = torch.randn(4, 4096, dtype=torch.bfloat16, device=DEVICE)
     gate = torch.randn(4, 4096, dtype=torch.bfloat16, device=DEVICE)

test/registered/jit/test_sparse_mla_q8kv8_prefill_sm90.py

@@ -114,7 +114,7 @@ def _torch_sparse_attention_ref(
 
 
 def _run_and_check(d_qk, with_sink, s_q=2, topk=TOPK, s_kv=S_KV):
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 
@@ -212,7 +212,7 @@ def test_sparse_mla_q8kv8_prefill_precision(d_qk: int, s_q: int, topk: int, s_kv
     """Demonstrate that Q8KV8 kernel precision is near-lossless versus the
     fp32 reference: max/mean/p99 absolute error are small and the fraction
     of elements exceeding 0.1 absolute error is under 1%."""
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 
@@ -285,7 +285,7 @@ def test_sparse_mla_q8kv8_prefill_precision(d_qk: int, s_q: int, topk: int, s_kv
 def test_sparse_mla_q8kv8_prefill_no_alias_between_calls():
     """Two default-allocation calls with the same shape must return independent
     storage. This guards against regressing to a module-scope output cache."""
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 
@@ -332,7 +332,7 @@ def test_sparse_mla_q8kv8_prefill_no_alias_between_calls():
 def test_sparse_mla_q8kv8_prefill_caller_owned_buffers():
     """Caller-provided ``out`` / ``max_logits`` / ``lse`` tensors must be
     written into in-place and returned as-is."""
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 
@@ -373,7 +373,7 @@ def test_sparse_mla_q8kv8_prefill_caller_owned_buffers():
 )
 def test_sparse_mla_q8kv8_prefill_rejects_bad_buffers():
     """Validation: wrong shape/dtype and aliasing must raise ValueError."""
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 
@@ -474,7 +474,7 @@ def test_sparse_mla_q8kv8_prefill_masked_sentinels(s_q: int):
     Gate: per-band cos > 0.97 AND magnitude ratio > 0.9
     (the denominator-pollution bug crushes magnitude 50-2000x, unmistakable even
     under fp8 noise)."""
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 
@@ -532,7 +532,7 @@ def test_sparse_mla_q8kv8_prefill_sq_envelope(s_q: int):
     """NEW gate (bug class 2): all-valid correctness across the s_q envelope.
     s_q=6144 previously produced first-band NaNs (an is_kv_valid data race that
     only appears past s_q=2048)."""
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 
@@ -575,7 +575,7 @@ def test_sparse_mla_q8kv8_prefill_large_skv():
     """NEW gate (bug class 3): large gathered buffers / large index values
     (s_kv=65536, indices in [33000, 65536)).  E2E multi-request gather buffers
     reach tens of thousands of rows; the suite above used s_kv<=1024."""
-    from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import (
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
         sparse_mla_q8kv8_prefill_fwd,
     )
 

test/registered/jit/test_timestep_embedding.py

@@ -10,10 +10,10 @@ try:
 except Exception:
     tabulate = None
 
-from sglang.jit_kernel.timestep_embedding import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.diffusion.timestep_embedding import (
     timestep_embedding as timestep_embedding_cuda,
 )
-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/kernels/test_cute_dsl_fp8_paged_mqa_logits.py

@@ -6,7 +6,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.dsa import cutedsl_paged_mqa_logits, pick_dsl_expand
+from sglang.kernels.ops.attention.dsa import cutedsl_paged_mqa_logits, pick_dsl_expand
 from sglang.srt.layers.attention.dsa.utils import (
     fp8_mqa_logits_ceil_to_ue8m0,
     fp8_mqa_logits_make_fused_kv,

test/registered/kernels/test_deepgemm_paged_mqa_logits.py

@@ -6,7 +6,7 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.dsa import (
+from sglang.kernels.ops.attention.dsa import (
     deepgemm_paged_mqa_logits_native,
     deepgemm_paged_mqa_logits_split,
 )

test/registered/kernels/test_dsa_indexer.py

@@ -643,7 +643,7 @@ class TestDSAIndexer(CustomTestCase):
         # preprocessed alongside the metadata (it asserts rather than silently
         # recomputing it) -- mirror what init_forward_metadata /
         # _build_forward_metadata_cuda_graph do.
-        from sglang.jit_kernel.dsv4.topk import plan_topk_v2
+        from sglang.kernels.ops.attention.dsv4.topk import plan_topk_v2
 
         attn_metadata = DSAMetadata(
             page_size=1,

test/registered/kernels/test_sigmoid_gate_mul.py

@@ -30,7 +30,7 @@ def reference_sigmoid_gate_mul(x, gate):
     ],
 )
 def test_sigmoid_gate_mul_correctness(shape, dtype):
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import sigmoid_gate_mul
 
     torch.manual_seed(42)
     x = torch.randn(shape, dtype=dtype, device=DEVICE)
@@ -60,7 +60,7 @@ def test_sigmoid_gate_mul_correctness(shape, dtype):
     ],
 )
 def test_sigmoid_gate_mul_broadcast_correctness(shape, dtype):
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import (
         sigmoid_gate_mul_broadcast,
     )
 
@@ -79,7 +79,7 @@ def test_sigmoid_gate_mul_broadcast_correctness(shape, dtype):
 
 @pytest.mark.parametrize("shape", [(4, 4096), (1, 128)])
 def test_sigmoid_gate_mul_broadcast_does_not_modify_inputs(shape):
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import (
         sigmoid_gate_mul_broadcast,
     )
 
@@ -97,7 +97,7 @@ def test_sigmoid_gate_mul_broadcast_does_not_modify_inputs(shape):
 
 
 def test_sigmoid_gate_mul_broadcast_output_dtype():
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import (
         sigmoid_gate_mul_broadcast,
     )
 
@@ -109,7 +109,7 @@ def test_sigmoid_gate_mul_broadcast_output_dtype():
 
 
 def test_sigmoid_gate_mul_broadcast_contiguous_output():
-    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+    from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import (
         sigmoid_gate_mul_broadcast,
     )
 

test/registered/kv_canary/test_self_unit_buffer_alloc.py

@@ -5,7 +5,7 @@ import unittest
 
 import torch
 
-from sglang.jit_kernel.kv_canary.consts import RealKvHashMode
+from sglang.kernels.ops.kv_canary.consts import RealKvHashMode
 from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
 from sglang.srt.kv_canary.pool_patcher.buffer_alloc import (
     make_packed_source,

test/registered/kv_canary/test_self_unit_endpoint.py

@@ -6,13 +6,13 @@ from unittest.mock import patch
 
 import torch
 
-from sglang.jit_kernel.kv_canary.consts import RealKvHashMode
-from sglang.jit_kernel.kv_canary.verify import (
+from sglang.kernels.ops.kv_canary.consts import RealKvHashMode
+from sglang.kernels.ops.kv_canary.verify import (
     CANARY_SLOT_BYTES,
     CanaryLaunchTag,
     VerifyPlan,
 )
-from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.kernels.ops.kv_canary.write import WritePlan
 from sglang.srt.kv_canary import endpoint as endpoint_module
 from sglang.srt.kv_canary.endpoint import (
     CanaryEndpoint,

test/registered/kv_canary/test_self_unit_perturb.py

@@ -7,7 +7,7 @@ from unittest.mock import patch
 
 import torch
 
-from sglang.jit_kernel.kv_canary.verify import RealKvSource
+from sglang.kernels.ops.kv_canary.verify import RealKvSource
 from sglang.srt.kv_canary.buffer_group import PoolKind
 from sglang.srt.kv_canary.perturb import (
     real_kv_post_forward,

test/registered/kv_canary/test_self_unit_pool_patcher.py

@@ -4,8 +4,8 @@ import unittest
 
 import torch
 
-from sglang.jit_kernel.kv_canary.consts import MAX_REAL_KV_SOURCES, RealKvHashMode
-from sglang.jit_kernel.kv_canary.verify import (
+from sglang.kernels.ops.kv_canary.consts import MAX_REAL_KV_SOURCES, RealKvHashMode
+from sglang.kernels.ops.kv_canary.verify import (
     CANARY_SLOT_BYTES,
     CanaryLaunchTag,
     RealKvSource,

test/registered/kv_canary/test_self_unit_runner_health.py

@@ -6,7 +6,7 @@ from unittest.mock import Mock
 
 import torch
 
-from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag
 from sglang.srt.kv_canary.config import CanaryConfig
 from sglang.srt.kv_canary.runner import stats_logger as stats_logger_module
 from sglang.srt.kv_canary.runner.health_checker import KernelRunCounterHealthChecker

test/registered/kv_canary/test_self_unit_runner_per_forward.py

@@ -5,9 +5,9 @@ from unittest.mock import patch
 
 import torch
 
-from sglang.jit_kernel.kv_canary.consts import RealKvHashMode
-from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan
-from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.kernels.ops.kv_canary.consts import RealKvHashMode
+from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag, VerifyPlan
+from sglang.kernels.ops.kv_canary.write import WritePlan
 from sglang.srt.kv_canary import endpoint as endpoint_module
 from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
 from sglang.srt.kv_canary.runner import kernel_launcher as kernel_launcher_module

test/registered/kv_canary/test_self_unit_runner_swa_divergence.py

@@ -6,7 +6,7 @@ from types import SimpleNamespace
 
 import torch
 
-from sglang.jit_kernel.kv_canary.verify import VerifyPlan
+from sglang.kernels.ops.kv_canary.verify import VerifyPlan
 from sglang.srt.environ import envs
 from sglang.srt.kv_canary.buffer_group import PoolKind
 from sglang.srt.kv_canary.runner import swa_divergence as swa_div_module

test/registered/kv_canary/test_self_unit_violation.py

@@ -6,9 +6,9 @@ from unittest.mock import patch
 
 import torch
 
-from sglang.jit_kernel.kv_canary import consts
-from sglang.jit_kernel.kv_canary.consts import FailReason
-from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.consts import FailReason
+from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag
 from sglang.srt.kv_canary.config import CanaryMode
 from sglang.srt.kv_canary.runner import violation_reporter as violation_reporter_module
 from sglang.srt.kv_canary.runner.violation_reporter import (

test/registered/lora/test_fused_moe_lora_kernel.py

@@ -5,11 +5,12 @@ import sys
 import pytest
 import torch
 
+from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora
+
 # ==============================================================================
 # IMPORT PREBUILT KERNEL
 # ==============================================================================
-from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
-from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora
+from sglang.kernels.ops.moe.moe_lora_align import moe_lora_align_block_size
 from sglang.srt.utils import set_random_seed
 from sglang.test.ci.ci_register import register_cuda_ci
 

test/registered/mock_model/test_self_unit_oracle.py

@@ -5,7 +5,7 @@ import unittest
 
 import torch
 
-from sglang.jit_kernel.kv_canary.consts import splitmix64
+from sglang.kernels.ops.kv_canary.consts import splitmix64
 from sglang.srt.kv_canary.token_oracle.oracle import (
     HashOracle,
     _splitmix64_tensor,

test/registered/mock_model/test_self_unit_oracle_torch_vs_ref.py

@@ -5,7 +5,7 @@ import unittest
 
 import torch
 
-from sglang.jit_kernel.kv_canary.verify_ref import splitmix64
+from sglang.kernels.ops.kv_canary.verify_ref import splitmix64
 from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

test/registered/unit/lora/test_experimental_sgl_marlin_alignment.py

@@ -34,7 +34,7 @@ def _load_align_function():
     namespace = {
         "torch": torch,
         "jit_moe_align_block_size": sys.modules[
-            "sglang.jit_kernel.moe_align"
+            "sglang.kernels.ops.moe.moe_align"
         ].moe_align_block_size,
     }
     exec(compile(module, str(ALIGN_PATH), "exec"), namespace)
@@ -62,7 +62,7 @@ def test_experimental_alignment_geometry_and_empty_input(monkeypatch):
 
     monkeypatch.setitem(
         sys.modules,
-        "sglang.jit_kernel.moe_align",
+        "sglang.kernels.ops.moe.moe_align",
         types.SimpleNamespace(moe_align_block_size=fake_jit_align),
     )
     align = _load_align_function()
@@ -155,7 +155,7 @@ def _assert_shared_outer_merged_align_semantics(
 def test_multi_slot_shared_outer_merged_align_cuda_graph_parity(num_slots):
     """The fused hot path must replay with current multi-LoRA routing data."""
 
-    from sglang.jit_kernel.trtllm_lora_temp.moe_lora_merged_align import (
+    from sglang.kernels.ops.moe.trtllm_lora_temp.moe_lora_merged_align import (
         moe_lora_merged_align,
     )