[Kernel] RFC #29630 finale: retire sglang.jit_kernel into sglang.kernels (#32072)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

test/README.md

@@ -72,7 +72,7 @@ Parameters: `est_time` (seconds), `stage` + `runner_config` (target stage and ru
 
 Keep `est_time`, `stage`, `runner_config` as **literal values** — `run_suite.py` collects them by AST parsing.
 
-JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/jit_kernel/` and are imported by absolute path):
+JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/kernels/jit/` and are imported by absolute path):
 - Correctness tests: `test/registered/jit/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
 - Benchmarks: `test/registered/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`
 

test/registered/attention/unittests/KNOWN_FAILURES.md

@@ -45,7 +45,7 @@ ImportError: cannot import name 'flash_attn_varlen_func' from 'flash_attn'
 ```
 
 **Root cause**: `DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func`
-via `sglang.jit_kernel.flash_attention`. On SM 8.x / 9.x that resolves to
+via `sglang.kernels.ops.attention.flash_attention`. On SM 8.x / 9.x that resolves to
 sgl-kernel's FA3 build (works on H200). On other SMs, the JIT kernel falls
 back to the upstream `flash_attn` (FA2) wheel — but the
 `lmsysorg/sglang:nightly-dev-cu13` container's `flash_attn` package on

test/registered/attention/unittests/dsa/README.md

@@ -29,7 +29,7 @@ hardware/SDK. The variant tests live in `test_dsa.py` as
 |---|---|---|---|---|
 | `flashmla_sparse` | ✓ | ✓ | ✓ | SM>=9.0 + `sgl_kernel.flash_mla` |
 | `flashmla_kv` | ✓ | ✓ | ✓ | SM>=9.0 + `sgl_kernel.flash_mla` |
-| `fa3` | ✓ | ✓ | ✓ | SM>=9.0 + `sglang.jit_kernel.flash_attention` |
+| `fa3` | ✓ | ✓ | ✓ | SM>=9.0 + `sglang.kernels.ops.attention.flash_attention` |
 | `tilelang` | ✓ (topk=2048 dedicated fixture) | ✓ (topk=2048 dedicated fixture) | skipped: not yet wired into CG runner | `tilelang_sparse_fwd` asserts `topk == 2048`; the topk=2048 fixture instance (`build_dsa_sparse_attention_fixture(..., index_topk=2048)`) is used by `test_sparse_tilelang_prefill_case` / `test_sparse_tilelang_decode_case`. The default-topk impl-variant matrix still skips tilelang with the same reason. **SM10.x container gate**: on Blackwell the tilelang JIT generates `wait_wgmma` which the container's MMA template doesn't ship (`KNOWN_FAILURES.md §2`); `dsa_impl_capability("tilelang")` skips on `major >= 10`. Set `SGLANG_TEST_DSA_TILELANG_FORCE=1` to override after re-imaging. |
 | `trtllm` | skipped: SM<10 | skipped: SM<10 | skipped: SM<10 | TRT-LLM Gen FMHA/MLA requires Blackwell (SM>=10.0). |
 | `aiter` | skipped: not HIP | skipped: not HIP | skipped: not HIP | AMD-only kernel library. |

test/registered/attention/unittests/dual_chunk/README.md

@@ -46,7 +46,7 @@ Columns are runner modes; rows are kernel-path modes of the single
 ## Container Gate (SM10.x)
 
 `DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func` via
-`sglang.jit_kernel.flash_attention`. On SM8.x / SM9.x that resolves to sgl-kernel's
+`sglang.kernels.ops.attention.flash_attention`. On SM8.x / SM9.x that resolves to sgl-kernel's
 FA3 build; on SM != {8, 9} (notably SM10.x / GB300) the JIT kernel falls back
 to the upstream `flash_attn` (FA2) wheel, which the
 `lmsysorg/sglang:nightly-dev-cu13` container ships without an SM10.x-compiled

test/registered/jit/benchmark/bench_activation.py

@@ -4,12 +4,18 @@ from sgl_kernel import gelu_and_mul as gelu_and_mul_aot
 from sgl_kernel import gelu_tanh_and_mul as gelu_tanh_and_mul_aot
 from sgl_kernel import silu_and_mul as silu_and_mul_aot
 
-from sglang.jit_kernel.activation import gelu_and_mul as gelu_and_mul_jit
-from sglang.jit_kernel.activation import gelu_tanh_and_mul as gelu_tanh_and_mul_jit
-from sglang.jit_kernel.activation import relu2 as relu2_jit
-from sglang.jit_kernel.activation import silu_and_mul as silu_and_mul_jit
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import create_random
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_random
+from sglang.kernels.ops.activation._jit_activation import (
+    gelu_and_mul as gelu_and_mul_jit,
+)
+from sglang.kernels.ops.activation._jit_activation import (
+    gelu_tanh_and_mul as gelu_tanh_and_mul_jit,
+)
+from sglang.kernels.ops.activation._jit_activation import relu2 as relu2_jit
+from sglang.kernels.ops.activation._jit_activation import (
+    silu_and_mul as silu_and_mul_jit,
+)
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_add_constant.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark_no_cudagraph,

test/registered/jit/benchmark/bench_awq_dequantize.py

@@ -4,7 +4,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import run_benchmark
+from sglang.kernels.jit.benchmark.utils import run_benchmark
 from sglang.kernels.ops.quantization.awq_dequantize import (
     awq_dequantize as jit_awq_dequantize,
 )

test/registered/jit/benchmark/bench_clamp_position.py

@@ -4,7 +4,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark,

test/registered/jit/benchmark/bench_concat_mla.py

@@ -6,7 +6,7 @@ import triton.testing
 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.kernels.jit.benchmark.utils import run_benchmark
 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

test/registered/jit/benchmark/bench_custom_all_reduce.py

@@ -10,8 +10,8 @@ import torch
 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.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range, multigpu_bench_main
 from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl
 from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py

@@ -8,9 +8,9 @@ import torch
 import torch.nn.functional as F
 import triton.testing
 
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
+from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm import dsv3_fused_a_gemm
 from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import (
     dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm,
 )

test/registered/jit/benchmark/bench_dsv3_router_gemm.py

@@ -1,16 +1,16 @@
 """Benchmark for DeepSeek V3 router GEMM (JIT kernel vs torch).
 
 Run on a Hopper (SM90+) GPU:
-    python -m sglang.jit_kernel.benchmark.bench_dsv3_router_gemm
+    python -m sglang.kernels.jit.benchmark.bench_dsv3_router_gemm
 """
 
 import torch
 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.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_random
 from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.kernels.ops.gemm._jit_dsv3_router_gemm import dsv3_router_gemm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py

@@ -6,7 +6,7 @@ import torch
 import triton
 
 from sglang.benchmark.bench_utils import run_bench
-from sglang.jit_kernel.benchmark.utils import get_benchmark_range
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range
 from sglang.srt.utils import is_sm100_supported
 from sglang.test.ci.ci_register import register_cuda_ci
 

test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py

@@ -5,7 +5,7 @@ import sys
 import torch
 import triton
 
-from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.kernels.ops.gemm.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
 from sglang.srt.utils import is_sm120_supported
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

test/registered/jit/benchmark/bench_fused_eh_norm.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import torch
 
-from sglang.jit_kernel.benchmark import marker
+from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py

@@ -1,6 +1,6 @@
 import torch
 
-from sglang.jit_kernel.benchmark import marker
+from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.ops.kvcache.fused_fp8_qkv_kv_cache import fused_fp8_qkv_kv_cache
 from sglang.test.ci.ci_register import register_cuda_ci
 

test/registered/jit/benchmark/bench_fused_qknorm_rope.py

@@ -14,7 +14,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.kernels.ops.attention.fused_qknorm_rope import (
     fused_qk_norm_rope as fused_qk_norm_rope_jit,
 )

test/registered/jit/benchmark/bench_hadamard.py

@@ -7,7 +7,7 @@ import torch.nn.functional as F
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     DEFAULT_DTYPE,
     get_benchmark_range,

test/registered/jit/benchmark/bench_hicache.py

@@ -25,7 +25,7 @@ import triton
 import triton.testing
 from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
 
-from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range
+from sglang.kernels.jit.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range
 from sglang.kernels.ops.kvcache.hicache import (
     can_use_hicache_jit_kernel,
     transfer_hicache_all_layer,

test/registered/jit/benchmark/bench_hisparse.py

@@ -5,7 +5,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
+from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
 from sglang.kernels.ops.kvcache.hisparse import load_cache_to_device_buffer_mla
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py

@@ -8,7 +8,7 @@ import triton
 import triton.language as tl
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     DEFAULT_DTYPE,
     DEFAULT_QUANTILES,

test/registered/jit/benchmark/bench_moe_fused_gate.py

@@ -1,7 +1,7 @@
 import torch
 
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import create_random
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_random
 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
 

test/registered/jit/benchmark/bench_mxfp8_moe.py

@@ -6,7 +6,7 @@ from typing import Any
 import torch
 import triton
 
-from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.kernels.ops.quantization.mxfp8 import (
     es_sm100_mxfp8_blockscaled_grouped_quant,
     es_sm100_mxfp8_blockscaled_moe_grouped_gemm,

test/registered/jit/benchmark/bench_ngram_compute_decode.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark_no_cudagraph,

test/registered/jit/benchmark/bench_ngram_update_token_table.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark_no_cudagraph,

test/registered/jit/benchmark/bench_norm.py

@@ -6,9 +6,11 @@ import triton.testing
 from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm
 from flashinfer.norm import rmsnorm as fi_rmsnorm
 
-from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
-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.benchmark.utils import get_benchmark_range, run_benchmark
+from sglang.kernels.ops.layernorm._jit_norm import (
+    fused_add_rmsnorm as jit_fused_add_rmsnorm,
+)
+from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_online_c128_mtp.py

@@ -9,7 +9,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark,

test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py

@@ -4,8 +4,10 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
-from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark
+from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import (
+    per_tensor_quant_fp8,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/jit/benchmark/bench_per_token_group_quant.py

@@ -1,10 +1,12 @@
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import create_empty, create_random
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_empty, create_random
 
 # per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
 # kernel is kept only as the perf baseline for this benchmark.
-from sglang.jit_kernel.per_token_group_quant import per_token_group_quant
-from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
+from sglang.kernels.ops.quantization._jit_per_token_group_quant import (
+    per_token_group_quant,
+)
+from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
     per_token_group_quant_8bit_v2,
 )
 from sglang.kernels.ops.quantization.fp8_kernel import (

test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py

@@ -1,9 +1,9 @@
 import torch
 from sgl_kernel import sgl_per_token_group_quant_8bit
 
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import create_random
-from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_random
+from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
     per_token_group_quant_8bit_v2,
 )
 from sglang.kernels.ops.quantization.fp8_kernel import (

test/registered/jit/benchmark/bench_per_token_group_quant_masked.py

@@ -2,13 +2,15 @@ import math
 
 import torch
 
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import create_empty, create_random
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_empty, create_random
 
 # per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
 # kernel is kept only as the perf baseline for this benchmark.
-from sglang.jit_kernel.per_token_group_quant import per_token_group_quant
-from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
+from sglang.kernels.ops.quantization._jit_per_token_group_quant import (
+    per_token_group_quant,
+)
+from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
     per_token_group_quant_8bit_v2,
 )
 from sglang.kernels.ops.quantization.fp8_kernel import (

test/registered/jit/benchmark/bench_post_reorder_deepgemm.py

@@ -1,6 +1,6 @@
 import torch
 
-from sglang.jit_kernel.benchmark import marker
+from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.ops.moe.ep_moe_kernels import (
     post_reorder_deepgemm,
     post_reorder_triton_kernel,

test/registered/jit/benchmark/bench_qknorm.py

@@ -1,8 +1,8 @@
 import torch
 
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import create_random
-from sglang.jit_kernel.norm import fused_inplace_qknorm
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_random
+from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm
 from sglang.srt.utils import get_current_device_stream_fast
 from sglang.test.ci.ci_register import register_cuda_ci
 

test/registered/jit/benchmark/bench_qknorm_across_heads.py

@@ -6,8 +6,8 @@ import triton
 import triton.testing
 from sgl_kernel import rmsnorm
 
-from sglang.jit_kernel.benchmark.utils import run_benchmark
-from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
+from sglang.kernels.jit.benchmark.utils import run_benchmark
+from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm_across_heads
 from sglang.srt.utils import get_current_device_stream_fast
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

test/registered/jit/benchmark/bench_renorm.py

@@ -5,7 +5,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
+from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/jit/benchmark/bench_resolve_future_token_ids.py

@@ -4,7 +4,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark,

test/registered/jit/benchmark/bench_rope.py

@@ -4,7 +4,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     DEFAULT_DTYPE,
     get_benchmark_range,

test/registered/jit/benchmark/bench_set_mla_kv_buffer.py

@@ -14,14 +14,16 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     DEFAULT_DTYPE,
     DEFAULT_QUANTILES,
     get_benchmark_range,
 )
-from sglang.jit_kernel.set_mla_kv_buffer import set_mla_kv_buffer as jit_set
 from sglang.kernels.jit.utils import is_arch_support_pdl
+from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import (
+    set_mla_kv_buffer as jit_set,
+)
 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_sparse_mla_q8kv8_prefill_sm90.py

@@ -6,7 +6,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
+from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph
 from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
     sparse_mla_q8kv8_prefill_fwd,
 )

test/registered/jit/benchmark/bench_spec_topk1.py

@@ -6,7 +6,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark,

test/registered/jit/benchmark/bench_store_cache.py

@@ -1,7 +1,7 @@
 import torch
 
-from sglang.jit_kernel.benchmark import marker
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     create_empty,
     create_random,

test/registered/jit/benchmark/bench_symm_mem_all_gather.py

@@ -26,8 +26,8 @@ import torch
 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.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range, multigpu_bench_main
 from sglang.kernels.jit.utils import cache_once
 from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
     all_gather_inner,

test/registered/jit/benchmark/bench_topk.py

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

test/registered/jit/benchmark/bench_tp_qknorm.py

@@ -23,8 +23,8 @@ import torch
 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 multigpu_bench_main
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import multigpu_bench_main
 from sglang.kernels.jit.utils import cache_once, get_ci_test_range
 from sglang.kernels.ops.communication.all_reduce import (
     fused_parallel_qknorm,

test/registered/jit/benchmark/bench_vocab_parallel_embedding.py

@@ -3,7 +3,7 @@ import torch.nn.functional as F
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
+from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.kernels.ops.embeddings.vocab_parallel_embedding import (
     vocab_parallel_embedding,
 )

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

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 
 import torch
 
-from sglang.jit_kernel.benchmark import marker
+from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.ops.diffusion.causal_conv3d_cat_pad import (
     fused_causal_conv3d_cat_pad_cuda,
 )

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

@@ -10,7 +10,7 @@ from typing import Any, Callable
 import flashinfer
 import torch
 
-from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE
+from sglang.kernels.jit.benchmark.utils import DEFAULT_DTYPE
 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
@@ -26,7 +26,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent
 REPO_ROOT = (
     Path(os.environ["SGLANG_NVFP4_REPO_ROOT"])
     if os.environ.get("SGLANG_NVFP4_REPO_ROOT")
-    # Anchor on the installed jit_kernel package (python/sglang/jit_kernel) so
+    # Anchor on the installed jit_kernel package (python/sglang/kernels/jit) so
     # this stays correct regardless of where the benchmark file lives.
     else KERNEL_PATH.parents[2]
 )

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

@@ -8,7 +8,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
+from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph
 from sglang.multimodal_gen.runtime.layers.layernorm import (
     LayerNormScaleShift,
     RMSNormScaleShift,

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

@@ -13,12 +13,14 @@ from typing import Callable
 import torch
 import torch.nn.functional as F
 
-from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE
-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.benchmark.utils import DEFAULT_DEVICE
 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.kernels.ops.layernorm._jit_norm import (
+    fused_add_rmsnorm as jit_fused_add_rmsnorm,
+)
+from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm
 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

@@ -5,7 +5,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     DEFAULT_DTYPE,
     get_benchmark_range,
@@ -131,7 +131,7 @@ def clone_inputs(
 def split_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
     from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
 
-    from sglang.jit_kernel.norm import fused_inplace_qknorm
+    from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm
 
     q = inputs["q"]
     k = inputs["k"]

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

@@ -3,7 +3,7 @@ from typing import Tuple
 import torch
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
+from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph
 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,

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

@@ -7,7 +7,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.kv_canary.utils import (
+from sglang.kernels.jit.benchmark.kv_canary.utils import (
     POOL_AXIS,
     SWA_WINDOW,
     BenchCase,
@@ -15,7 +15,7 @@ from sglang.jit_kernel.benchmark.kv_canary.utils import (
     build_full_matrix_cases,
     naive_cumsum_fn,
 )
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark,

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

@@ -6,7 +6,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark_no_cudagraph,

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

@@ -6,7 +6,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.kv_canary.utils import (
+from sglang.kernels.jit.benchmark.kv_canary.utils import (
     RING_CAPACITY,
     SWA_WINDOW,
     BenchCase,
@@ -16,7 +16,7 @@ from sglang.jit_kernel.benchmark.kv_canary.utils import (
     make_real_kv_sources,
     naive_slot_copy_fn,
 )
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark,

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

@@ -6,7 +6,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.kv_canary.utils import (
+from sglang.kernels.jit.benchmark.kv_canary.utils import (
     RING_CAPACITY,
     SWA_WINDOW,
     BenchCase,
@@ -16,7 +16,7 @@ from sglang.jit_kernel.benchmark.kv_canary.utils import (
     make_real_kv_sources,
     naive_slot_copy_fn,
 )
-from sglang.jit_kernel.benchmark.utils import (
+from sglang.kernels.jit.benchmark.utils import (
     DEFAULT_DEVICE,
     get_benchmark_range,
     run_benchmark,

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

@@ -9,7 +9,7 @@ intermediate buffers; the baseline is two launches with split-K partials.
 import torch
 import triton
 
-from sglang.jit_kernel.benchmark import marker
+from sglang.kernels.jit.benchmark import marker
 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,

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

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

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

@@ -3,7 +3,7 @@ per-buffer index_put_ stores (main K, main V, index K, optional index V)."""
 
 import torch
 
-from sglang.jit_kernel.benchmark import marker
+from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.ops.kvcache.minimax_store_kv_index import store_kv_index
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 

test/registered/jit/deepseek_v4/test_c128_v2.py

@@ -7,7 +7,7 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.tests.deepseek_v4.common import (
+from sglang.kernels.jit.tests.deepseek_v4.common import (
     LegacyContext,
     PagedContext,
     make_legacy_context,

test/registered/jit/deepseek_v4/test_c4_v2.py

@@ -7,7 +7,7 @@ import pytest
 import torch
 import triton
 
-from sglang.jit_kernel.tests.deepseek_v4.common import (
+from sglang.kernels.jit.tests.deepseek_v4.common import (
     LegacyContext,
     PagedContext,
     make_legacy_context,

test/registered/jit/diffusion/test_qknorm_rope.py

@@ -48,7 +48,7 @@ def split_qknorm_rope(
 ) -> None:
     from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
 
-    from sglang.jit_kernel.norm import fused_inplace_qknorm
+    from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm
 
     fused_inplace_qknorm(q, k, q_weight, k_weight)
     apply_rope_with_cos_sin_cache_inplace(

test/registered/jit/kv_canary/test_const_sync.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import re
 from pathlib import Path
 
-import sglang.jit_kernel
+import sglang.kernels.jit
 from sglang.kernels.ops.kv_canary import consts
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
@@ -11,10 +11,10 @@ register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-la
 register_amd_ci(est_time=5, stage="jit-kernel-unit", runner_config="amd")
 
 
-# Resolve the kernel source against the installed jit_kernel package rather
-# than this file's location, so the test stays correct wherever it lives.
+# Resolve the kernel source against the installed sglang.kernels.jit package
+# rather than this file's location, so the test stays correct wherever it lives.
 _CONSTS_CUH: Path = (
-    Path(sglang.jit_kernel.__file__).resolve().parent
+    Path(sglang.kernels.jit.__file__).resolve().parent
     / "csrc"
     / "kv_canary"
     / "consts.cuh"

test/registered/jit/kv_canary/test_kernel_config.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
+from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_canary_buf_equal,
     assert_canary_state_equal,
@@ -15,13 +15,13 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     make_write_plan_pair,
     stamp_clean_chain,
 )
-from sglang.jit_kernel.tests.kv_canary._differential import (
+from sglang.kernels.jit.tests.kv_canary._differential import (
     _assert_plans_byte_equal,
     _run_both_plan,
     _run_both_verify,
     _run_both_write,
 )
-from sglang.jit_kernel.tests.kv_canary._fixtures import (
+from sglang.kernels.jit.tests.kv_canary._fixtures import (
     dummy_pseudo_tensors,
     empty_extras,
 )

test/registered/jit/kv_canary/test_pipeline_e2e.py

@@ -5,7 +5,7 @@ from typing import Any, Optional
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
+from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_canary_buf_equal,
     assert_canary_state_equal,
@@ -14,7 +14,7 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     stamp_clean_chain,
     write_slot_fields,
 )
-from sglang.jit_kernel.tests.kv_canary._fixtures import (
+from sglang.kernels.jit.tests.kv_canary._fixtures import (
     clone_real_kv_sources,
     empty_extras,
     make_req_to_token,

test/registered/jit/kv_canary/test_plan_fuzz.py

@@ -7,19 +7,19 @@ from typing import Optional
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._differential import _run_both_plan
-from sglang.jit_kernel.tests.kv_canary._fixtures import (
+from sglang.kernels.jit.tests.kv_canary._differential import _run_both_plan
+from sglang.kernels.jit.tests.kv_canary._fixtures import (
     allocate_plan_pair,
     derive_plan_capacity,
     make_lut,
     make_padding_mask,
     make_req_to_token,
 )
-from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
+from sglang.kernels.jit.tests.kv_canary._fuzz_driver import (
     FUZZ_SEEDS_PR,
     run_fuzz_combo,
 )
-from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
+from sglang.kernels.jit.tests.kv_canary._invariants import PlanInvariants
 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,15 +5,15 @@ import random
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._differential import run_plan_diff
-from sglang.jit_kernel.tests.kv_canary._fixtures import (
+from sglang.kernels.jit.tests.kv_canary._differential import run_plan_diff
+from sglang.kernels.jit.tests.kv_canary._fixtures import (
     allocate_plan_pair,
     derive_plan_capacity,
     empty_extras,
     make_lut,
     make_req_to_token,
 )
-from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
+from sglang.kernels.jit.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,

test/registered/jit/kv_canary/test_utils.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from sglang.jit_kernel.benchmark.kv_canary.utils import (
+from sglang.kernels.jit.benchmark.kv_canary.utils import (
     MAX_EXTEND_TOKENS_PER_FORWARD,
     build_fast_matrix_cases,
     build_full_matrix_cases,

test/registered/jit/kv_canary/test_verify_fuzz.py

@@ -6,23 +6,23 @@ from dataclasses import dataclass
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
+from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     make_canary_buf,
     make_log_pair,
     make_verify_plan_pair,
     stamp_clean_chain,
 )
-from sglang.jit_kernel.tests.kv_canary._differential import _run_both_verify
-from sglang.jit_kernel.tests.kv_canary._fixtures import (
+from sglang.kernels.jit.tests.kv_canary._differential import _run_both_verify
+from sglang.kernels.jit.tests.kv_canary._fixtures import (
     clone_real_kv_sources,
     make_real_kv_sources,
 )
-from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
+from sglang.kernels.jit.tests.kv_canary._fuzz_driver import (
     FUZZ_SEEDS_PR,
     run_fuzz_combo,
 )
-from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants
+from sglang.kernels.jit.tests.kv_canary._invariants import VerifyInvariants
 from sglang.kernels.ops.kv_canary import consts
 from sglang.kernels.ops.kv_canary.verify import (
     CanaryLaunchTag,

test/registered/jit/kv_canary/test_verify_hand.py

@@ -8,7 +8,7 @@ from typing import Callable
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
+from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_only_bits_set,
     chain_anchor_signed,
@@ -26,12 +26,12 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     to_signed_int64,
     write_slot_fields,
 )
-from sglang.jit_kernel.tests.kv_canary._differential import (
+from sglang.kernels.jit.tests.kv_canary._differential import (
     _run_both_verify,
     run_verify_diff,
 )
-from sglang.jit_kernel.tests.kv_canary._fixtures import clone_real_kv_sources
-from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
+from sglang.kernels.jit.tests.kv_canary._fixtures import clone_real_kv_sources
+from sglang.kernels.jit.tests.kv_canary._hand_oracle import (
     _hand_fold_all,
     _hand_fold_partial,
 )

test/registered/jit/kv_canary/test_write_fuzz.py

@@ -6,23 +6,23 @@ from dataclasses import dataclass
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
+from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     make_canary_buf,
     make_log_pair,
     make_write_plan_pair,
     stamp_pair,
 )
-from sglang.jit_kernel.tests.kv_canary._differential import _run_both_write
-from sglang.jit_kernel.tests.kv_canary._fixtures import (
+from sglang.kernels.jit.tests.kv_canary._differential import _run_both_write
+from sglang.kernels.jit.tests.kv_canary._fixtures import (
     clone_real_kv_sources,
     make_real_kv_sources,
 )
-from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
+from sglang.kernels.jit.tests.kv_canary._fuzz_driver import (
     FUZZ_SEEDS_PR,
     run_fuzz_combo,
 )
-from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants
+from sglang.kernels.jit.tests.kv_canary._invariants import WriteInvariants
 from sglang.kernels.ops.kv_canary import consts
 from sglang.kernels.ops.kv_canary.verify import (
     CanaryLaunchTag,

test/registered/jit/kv_canary/test_write_hand.py

@@ -7,7 +7,7 @@ from unittest.mock import patch
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
+from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
     FakeViolationLog,
     assert_canary_state_equal,
     assert_only_bits_set,
@@ -24,15 +24,15 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
     stamp_pair,
     to_signed_int64,
 )
-from sglang.jit_kernel.tests.kv_canary._differential import (
+from sglang.kernels.jit.tests.kv_canary._differential import (
     _run_both_write,
     run_write_diff,
 )
-from sglang.jit_kernel.tests.kv_canary._fixtures import (
+from sglang.kernels.jit.tests.kv_canary._fixtures import (
     clone_real_kv_sources,
     dummy_pseudo_tensors,
 )
-from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
+from sglang.kernels.jit.tests.kv_canary._hand_oracle import (
     _hand_fold_all,
     _hand_fold_partial,
 )

test/registered/jit/test_activation.py

@@ -4,12 +4,12 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from sglang.jit_kernel.activation import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.activation._jit_activation import (
     SUPPORTED_ACTIVATIONS,
     relu2,
     run_activation,
 )
-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

@@ -28,7 +28,7 @@ import torch
 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.kernels.jit.tests.utils import multigpu_pytest_main
 from sglang.kernels.jit.utils import cache_once, get_ci_test_range
 from sglang.kernels.ops.communication.all_reduce import (
     AllReduceAlgo,

test/registered/jit/test_cutedsl_gdn.py

@@ -13,7 +13,7 @@ try:
     import cutlass  # noqa: F401
     from cutlass.cute.runtime import from_dlpack
 
-    from sglang.jit_kernel import cutedsl_gdn
+    from sglang.kernels.ops.attention import cutedsl_gdn
 
     CUTEDSL_AVAILABLE = True
 except ImportError:

test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py

@@ -89,7 +89,7 @@ from typing import Callable, Iterable, Literal, Optional
 import pytest
 import torch
 
-from sglang.jit_kernel.tests.deepseek_v4.common import (
+from sglang.kernels.jit.tests.deepseek_v4.common import (
     make_legacy_context,
     to_seq_extend,
 )

test/registered/jit/test_dsv3_fused_a_gemm.py

@@ -6,12 +6,12 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
 from sglang.kernels.jit.utils import (
     get_ci_test_range,
     get_jit_cuda_arch,
     is_hip_runtime,
 )
+from sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm import dsv3_fused_a_gemm
 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

@@ -6,12 +6,12 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm
 from sglang.kernels.jit.utils import (
     get_ci_test_range,
     get_jit_cuda_arch,
     is_hip_runtime,
 )
+from sglang.kernels.ops.gemm._jit_dsv3_router_gemm import dsv3_router_gemm
 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

@@ -20,7 +20,7 @@ def sglang_jit_fused_add_rmsnorm(
     *,
     cast_x_before_out_mul: bool = False,
 ) -> None:
-    from sglang.jit_kernel.norm import fused_add_rmsnorm
+    from sglang.kernels.ops.layernorm._jit_norm import fused_add_rmsnorm
 
     fused_add_rmsnorm(
         input, residual, weight, eps, cast_x_before_out_mul=cast_x_before_out_mul

test/registered/jit/test_hadamard_jit.py

@@ -19,7 +19,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(est_time=128, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 # Exact M×N Hadamard matrices (±1 entries) copied from
-# python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py.
+# python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py.
 # These are non-power-of-2 Hadamard matrices constructed via Paley/Williamson methods.
 # "+" = +1, "-" = -1.  Used by the _12n/_20n/_28n/_40n kernel variants.
 
@@ -145,7 +145,7 @@ def _parse_hadamard_str(s):
 
 
 # Parsed M×M special Hadamard matrices, keyed by M (the "multiple").
-# Copied from python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py
+# Copied from python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py
 # (had_12_paley, had_20_will, had_28_will, had_40_tpal)
 _SPECIAL_MATRICES = {
     12: _parse_hadamard_str(_HAD_12_STR),
@@ -187,7 +187,7 @@ def hadamard_transform_mn_ref(x, multiple, scale=1.0):
     x = x.reshape(-1, dim)
 
     # The kernel requires dim % (4*M) == 0 (for vectorized memory access).
-    # See python/sglang/jit_kernel/hadamard.py: pad_multiple = 4 * 12 / 4 * 20 / etc.
+    # See python/sglang/kernels/ops/attention/hadamard.py: pad_multiple = 4 * 12 / 4 * 20 / etc.
     pad_multiple = 4 * multiple
     if dim % pad_multiple != 0:
         pad_size = pad_multiple - dim % pad_multiple
@@ -319,7 +319,7 @@ def test_hadamard_transform_scale_one(dtype):
 
 # Test dimensions for M×N variants: dim = M * N where N = 2^k.
 # M = 12/20/28/40 are the non-power-of-2 Hadamard sizes registered in
-# python/sglang/jit_kernel/hadamard.py (Hadamard12NKernel, ..., Hadamard40NKernel).
+# python/sglang/kernels/ops/attention/hadamard.py (Hadamard12NKernel, ..., Hadamard40NKernel).
 # range(2,9) gives N = 4,8,...,256 so dims cover a practical range.
 _12N_DIMS = [12 * (2**k) for k in range(2, 9)]  # 48, 96, ... , 3072
 _20N_DIMS = [20 * (2**k) for k in range(2, 9)]  # 80, 160, ... , 5120

test/registered/jit/test_per_tensor_quant_fp8.py

@@ -5,8 +5,10 @@ from typing import Optional, Tuple
 import pytest
 import torch
 
-from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import (
+    per_tensor_quant_fp8,
+)
 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

@@ -21,8 +21,10 @@ import itertools
 import pytest
 import torch
 
-from sglang.jit_kernel.per_token_group_quant import per_token_group_quant
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.quantization._jit_per_token_group_quant import (
+    per_token_group_quant,
+)
 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

@@ -3,10 +3,10 @@ import itertools
 import pytest
 import torch
 
-from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
     per_token_group_quant_8bit_v2,
 )
-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

@@ -34,7 +34,7 @@ def sglang_jit_qknorm(
     q_weight: torch.Tensor,
     k_weight: torch.Tensor,
 ) -> None:
-    from sglang.jit_kernel.norm import fused_inplace_qknorm
+    from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm
 
     fused_inplace_qknorm(q, k, q_weight, k_weight)
 

test/registered/jit/test_qknorm_across_heads.py

@@ -19,7 +19,7 @@ def sglang_jit_qknorm_across_heads(
     q_weight: torch.Tensor,
     k_weight: torch.Tensor,
 ) -> None:
-    from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
+    from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm_across_heads
 
     fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight)
 

test/registered/jit/test_rmsnorm.py

@@ -26,7 +26,7 @@ def sglang_jit_rmsnorm(
     output: torch.Tensor | None = None,
     eps: float = EPS,
 ) -> None:
-    from sglang.jit_kernel.norm import rmsnorm
+    from sglang.kernels.ops.layernorm._jit_norm import rmsnorm
 
     rmsnorm(input, weight, out=output, eps=eps)
 
@@ -127,7 +127,7 @@ def test_rmsnorm(
 
 @pytest.mark.parametrize("hidden_size", [64, 128, 256, 512, 8192, 8704, 16384])
 def test_rmsnorm_hidden_size_support(hidden_size: int) -> None:
-    from sglang.jit_kernel.norm import _is_supported_rmsnorm_hidden_size
+    from sglang.kernels.ops.layernorm._jit_norm import _is_supported_rmsnorm_hidden_size
 
     assert _is_supported_rmsnorm_hidden_size(hidden_size)
 
@@ -148,7 +148,7 @@ def test_rmsnorm_hidden_size_support(hidden_size: int) -> None:
     ],
 )
 def test_rmsnorm_kernel_dispatch(hidden_size: int, expected: str) -> None:
-    from sglang.jit_kernel.norm import _rmsnorm_kernel_class
+    from sglang.kernels.ops.layernorm._jit_norm import _rmsnorm_kernel_class
 
     assert _rmsnorm_kernel_class(hidden_size) == expected
 

test/registered/jit/test_set_mla_kv_buffer.py

@@ -3,11 +3,11 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.set_mla_kv_buffer import (
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import (
     can_use_set_mla_kv_buffer,
     set_mla_kv_buffer,
 )
-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_symm_mem_all_gather.py

@@ -26,7 +26,7 @@ import torch
 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.kernels.jit.tests.utils import multigpu_pytest_main
 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,

test/registered/jit/test_tp_qknorm.py

@@ -14,7 +14,7 @@ import torch.distributed as dist
 import triton
 
 import sglang.srt.distributed.parallel_state as ps
-from sglang.jit_kernel.tests.utils import multigpu_pytest_main
+from sglang.kernels.jit.tests.utils import multigpu_pytest_main
 from sglang.kernels.jit.utils import cache_once
 from sglang.kernels.ops.communication.all_reduce import (
     fused_parallel_qknorm,

test/registered/kernels/test_kernels_namespace.py

@@ -183,11 +183,12 @@ def test_platform_detect_does_not_raise():
 
 
 def test_import_stays_metadata_only():
-    # Importing the namespace must not pull in sgl_kernel / sglang.jit_kernel.
+    # Importing the namespace must not pull in the AOT backend (sgl_kernel) or
+    # the JIT compilation infra (sglang.kernels.jit), which import torch / nvcc.
     code = (
         "import sys, sglang.kernels.ops; "
         "print('DIRTY' if 'sgl_kernel' in sys.modules or any("
-        "m.startswith('sglang.jit_kernel') for m in sys.modules) else 'CLEAN')"
+        "m.startswith('sglang.kernels.jit') for m in sys.modules) else 'CLEAN')"
     )
     r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True)
     assert r.returncode == 0, r.stderr

test/registered/unit/distributed/test_vmm_utils.py

@@ -20,7 +20,7 @@ import torch
 import torch.distributed as dist
 from cuda.bindings import driver as drv
 
-from sglang.jit_kernel.tests.utils import multigpu_pytest_main
+from sglang.kernels.jit.tests.utils import multigpu_pytest_main
 from sglang.kernels.jit.utils import cache_once
 from sglang.srt.distributed.device_communicators.vmm_utils import (
     check_drv,