Refactor JIT kernel and expert-pack directory layout (#36704)

docs/docs/developer_guide/development_jit_kernel_guide.mdx

@@ -168,7 +168,7 @@ def add_constant(src: torch.Tensor, c: int):
 
 ### STEP 1: Write the C++ kernel
 
-Write your CUDA kernel in [kernels/jit/csrc/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/jit/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
+Write your CUDA kernel in [kernels/jit/csrc/elementwise/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/jit/csrc/elementwise/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
 
 ```cpp Example
 #include <sgl_kernel/tensor.h>   // For TensorMatcher, SymbolicSize, SymbolicDevice
@@ -231,7 +231,7 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
 ### STEP 2: Create Python Interfaces
 
 Next, expose the kernel through a Python wrapper.
-Create a new file at [kernels/ops/attention/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/ops/attention/add_constant.py) and expose the needed interfaces.
+Create a new file at [kernels/ops/elementwise/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/ops/elementwise/add_constant.py) and expose the needed interfaces.
 
 ```python Example
 from __future__ import annotations
@@ -251,7 +251,7 @@ def _jit_add_constant_module(constant: int) -> Module:
     return load_jit(
         "add_constant",
         *args,
-        cuda_files=["add_constant.cuh"],
+        cuda_files=["elementwise/add_constant.cuh"],
         cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
     )
 
@@ -275,10 +275,10 @@ Keep the Python wrapper thin, but still validate the basic invariants such as de
 Finally, import and use the kernel like a regular Python function:
 
 ```python Example
-from sglang.kernels.jit.add_constant import add_constant
+from sglang.kernels.ops.elementwise.add_constant import add_constant
 ```
 
-For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/test/registered/jit/test_add_constant.py).
+For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/test/registered/kernels/ops/elementwise/test_add_constant.py).
 
 ## C++ Include Library Reference
 

examples/runtime/deepseek_v4/benchmark_deepseek_5090.py

@@ -44,7 +44,7 @@ def find_sglang_repo() -> Path:
         return Path(configured).expanduser().resolve()
     for candidate in (SCRIPT_DIR, *SCRIPT_DIR.parents):
         if (candidate / "python" / "sglang").is_dir() and (
-            candidate / "tools" / "expert_pack"
+            candidate / "python" / "sglang" / "srt" / "model_loader" / "expert_pack"
         ).is_dir():
             return candidate
     raise RuntimeError("could not locate the SGLang repository")

examples/runtime/kimi_k3/benchmark_kimi_k3_5090.py

@@ -45,7 +45,7 @@ def find_sglang_repo() -> Path:
         return Path(configured).expanduser().resolve()
     for candidate in (SCRIPT_DIR, *SCRIPT_DIR.parents):
         if (candidate / "python" / "sglang").is_dir() and (
-            candidate / "tools" / "expert_pack"
+            candidate / "python" / "sglang" / "srt" / "model_loader" / "expert_pack"
         ).is_dir():
             return candidate
     raise RuntimeError("could not locate the SGLang repository")

python/sglang/kernels/jit/csrc/elementwise/add_constant.cuh

@@ -1,3 +1,5 @@
+// SPDX-License-Identifier: Apache-2.0
+
 #include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
 #include <sgl_kernel/utils.h>   // For div_ceil, RuntimeCheck
 

python/sglang/kernels/jit/csrc/speculative/ngram_embedding.cuh

@@ -1,3 +1,5 @@
+// SPDX-License-Identifier: Apache-2.0
+
 #include <sgl_kernel/tensor.h>
 #include <sgl_kernel/utils.h>
 

python/sglang/kernels/jit/minicpm_sala/__init__.py

@@ -1,3 +0,0 @@
-from sglang.kernels.jit.minicpm_sala.get_block_table import get_block_table
-
-__all__ = ["get_block_table"]

python/sglang/kernels/ops/__init__.py

@@ -28,6 +28,7 @@ _GROUPS = (
     "layernorm",
     "mamba",
     "memory",
+    "minicpm_sala",
     "mm",
     "moe",
     "quantization",

python/sglang/kernels/ops/elementwise/add_constant.py

@@ -16,7 +16,7 @@ def _jit_add_constant_module(constant: int) -> Module:
     return load_jit(
         "add_constant",
         *args,
-        cuda_files=["add_constant.cuh"],
+        cuda_files=["elementwise/add_constant.cuh"],
         cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
     )
 

python/sglang/kernels/ops/minicpm_sala/__init__.py

@@ -0,0 +1,24 @@
+"""MiniCPM-SALA kernels."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from sglang.kernels.ops.minicpm_sala.get_block_table import get_block_table
+
+
+def __getattr__(name: str) -> Any:
+    if name == "get_block_table":
+        from sglang.kernels.ops.minicpm_sala.get_block_table import get_block_table
+
+        globals()[name] = get_block_table
+        return get_block_table
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(__all__))
+
+
+__all__ = ["get_block_table"]

python/sglang/kernels/ops/minicpm_sala/get_block_table.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

python/sglang/kernels/ops/speculative/ngram_embedding.py

@@ -14,7 +14,7 @@ if TYPE_CHECKING:
 def _jit_ngram_embedding_module() -> Module:
     return load_jit(
         "ngram_embedding",
-        cuda_files=["ngram_embedding.cuh"],
+        cuda_files=["speculative/ngram_embedding.cuh"],
         cuda_wrappers=[
             ("compute_n_gram_ids", "&NgramEmbeddingKernel::compute_n_gram_ids"),
             (

python/sglang/srt/layers/attention/minicpm/backend.py

@@ -25,7 +25,7 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
 
 
-from sglang.kernels.jit.minicpm_sala import get_block_table
+from sglang.kernels.ops.minicpm_sala import get_block_table
 from sglang.srt.layers.attention.minicpm.sparse_utils import (
     CompressionLevelMetadata,
     MiniCPMSparseMetadata,

python/sglang/srt/model_loader/expert_pack/__init__.py

@@ -0,0 +1 @@
+"""Expert-pack build, inspection, and validation utilities."""

python/sglang/srt/model_loader/expert_pack/build.py

@@ -215,14 +215,16 @@ def tool_sha256() -> str:
 
 
 def git_sha() -> str:
-    result = subprocess.run(
-        ["git", "rev-parse", "HEAD"],
-        cwd=Path(__file__).resolve().parents[2],
-        check=True,
-        capture_output=True,
-        text=True,
-    )
-    return result.stdout.strip()
+    try:
+        return subprocess.run(
+            ["git", "rev-parse", "HEAD"],
+            cwd=Path(__file__).resolve().parent,
+            check=True,
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+    except (OSError, subprocess.CalledProcessError):
+        return "unknown"
 
 
 def build(args: argparse.Namespace) -> dict[str, object]:

python/sglang/srt/model_loader/expert_pack/format.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
 from __future__ import annotations
 
 import hashlib

python/sglang/srt/model_loader/expert_pack/kimi_ggml.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
-"""Structural inventory and adapter manifest for Kimi K3 GGUF assets."""
+"""Structural inventory and adapter manifest for Kimi K3 GGUF expert packs."""
 
 from __future__ import annotations
 

python/sglang/srt/model_loader/expert_pack/prepare_deepseek_pack.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
 """Validate or build the DeepSeek expert-pack used by the RTX 5090 benchmark."""
 
 from __future__ import annotations

python/sglang/srt/model_loader/expert_pack/prepare_kimi_manifest.py

@@ -37,7 +37,7 @@ def main() -> int:
     args = parse_args()
     if args.payload_samples < 0:
         raise ValueError("--payload-samples must be non-negative")
-    repo = Path(__file__).resolve().parents[2]
+    repo = Path(__file__).resolve().parent
     manifest = create_manifest(
         gguf_dir=args.gguf_dir,
         expert_pack=args.expert_pack,

python/sglang/srt/model_loader/expert_pack/prepare_kimi_pack.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
-"""Validate or build the Kimi K3 Expert Pack derived from GGUF shards."""
+"""Validate or build the Kimi K3 expert pack derived from GGUF shards."""
 
 from __future__ import annotations
 

python/sglang/srt/model_loader/expert_pack/validate.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
 from __future__ import annotations
 
 import argparse

python/sglang/srt/model_loader/expert_pack_runtime.py

@@ -114,14 +114,19 @@ def _expert_pack_path(gguf: Path) -> Path:
     return gguf.parent / f"{gguf.name[: match.start()]}.expert-major.pack"
 
 
-def _repo_root() -> Path:
-    for candidate in Path(__file__).resolve().parents:
-        if (candidate / "tools" / "expert_pack" / "prepare_kimi_pack.py").is_file():
-            return candidate
-    raise RuntimeError(
-        "expert_pack cannot auto-build Kimi artifacts from an installed package; "
-        "run from an SGLang source checkout"
+def _expert_pack_tools_dir() -> Path:
+    tools_dir = Path(__file__).with_name("expert_pack")
+    required_tools = (
+        "prepare_deepseek_pack.py",
+        "prepare_kimi_manifest.py",
+        "prepare_kimi_pack.py",
     )
+    missing = [name for name in required_tools if not (tools_dir / name).is_file()]
+    if missing:
+        raise RuntimeError(
+            "expert_pack preparation tools are missing: " + ", ".join(missing)
+        )
+    return tools_dir
 
 
 def ensure_kimi_assets(
@@ -144,26 +149,26 @@ def ensure_kimi_assets(
     manifest = artifact_dir / "kimi-k3-expert-pack.manifest.json"
     tokenizer = resolve_kimi_tokenizer(gguf, tokenizer_dir)
     lock_path = pack.with_name(pack.name + ".startup.lock")
-    repo = _repo_root()
+    tools_dir = _expert_pack_tools_dir()
     with lock_path.open("w") as lock:
         fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
         model_dir = prepare_kimi_model_metadata(tokenizer, artifact_dir)
         subprocess.run(
             [
                 sys.executable,
-                str(repo / "tools" / "expert_pack" / "prepare_kimi_pack.py"),
+                str(tools_dir / "prepare_kimi_pack.py"),
                 "--gguf",
                 str(gguf),
                 "--model-config",
                 str(model_dir / "config.json"),
             ],
-            cwd=repo,
+            cwd=tools_dir,
             check=True,
         )
         subprocess.run(
             [
                 sys.executable,
-                str(repo / "tools" / "expert_pack" / "prepare_kimi_manifest.py"),
+                str(tools_dir / "prepare_kimi_manifest.py"),
                 "--gguf-dir",
                 str(gguf_dir),
                 "--expert-pack",
@@ -177,7 +182,7 @@ def ensure_kimi_assets(
                 "--payload-samples",
                 "6",
             ],
-            cwd=repo,
+            cwd=tools_dir,
             check=True,
         )
     return {
@@ -480,9 +485,9 @@ def _deepseek_digest(value: object, field: str) -> str:
 
 
 def _prepare_deepseek_pack(
-    source: Path, model_config: Path, repo: Path
+    source: Path, model_config: Path, tools_dir: Path
 ) -> tuple[Path, Path]:
-    tool = repo / "tools" / "expert_pack" / "prepare_deepseek_pack.py"
+    tool = tools_dir / "prepare_deepseek_pack.py"
     if not tool.is_file():
         raise FileNotFoundError(f"missing DeepSeek Expert Pack preparer: {tool}")
     subprocess.run(
@@ -494,7 +499,7 @@ def _prepare_deepseek_pack(
             "--model-config",
             str(model_config),
         ],
-        cwd=repo,
+        cwd=tools_dir,
         check=True,
     )
     return (
@@ -514,14 +519,14 @@ def prepare_raw_deepseek_server_args(
     source = Path(cfg.model_path).expanduser().resolve(strict=True)
     if not source.is_file():
         return
-    repo = _repo_root()
+    tools_dir = _expert_pack_tools_dir()
     artifact_dir = _deepseek_artifact_dir_for_source(source).resolve()
     lock_path = artifact_dir / "deepseek-v4-startup.lock"
     artifact_dir.mkdir(parents=True, exist_ok=True)
     with lock_path.open("w") as lock:
         fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
         model_config = _prepare_deepseek_model_metadata(source, artifact_dir)
-        pack, manifest = _prepare_deepseek_pack(source, model_config, repo)
+        pack, manifest = _prepare_deepseek_pack(source, model_config, tools_dir)
     manifest_value = json.loads(manifest.read_text(encoding="utf-8"))
     source_value = manifest_value.get("source") or {}
     model_value = manifest_value.get("model") or {}

test/registered/expert_pack/test_expert_pack_runtime.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 import hashlib
 import json
-import sys
 import tempfile
 import unittest
 from dataclasses import replace
@@ -16,11 +15,11 @@ from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=15, suite="base-a-test-cpu")
 
-ROOT = Path(__file__).resolve().parents[3]
-TOOLS = ROOT / "tools" / "expert_pack"
-sys.path.insert(0, str(TOOLS))
-
-from format import (  # noqa: E402
+from sglang.srt.layers.moe.expert_pack import (  # noqa: E402
+    ExpertPackStore,
+    _CacheSlot,
+)
+from sglang.srt.model_loader.expert_pack.format import (  # noqa: E402
     ENTRY_STRUCT,
     FLAG_IDENTITY_PAYLOAD,
     FLAG_TRIPLET_OBJECTS,
@@ -32,11 +31,6 @@ from format import (  # noqa: E402
     read_index,
 )
 
-from sglang.srt.layers.moe.expert_pack import (  # noqa: E402
-    ExpertPackStore,
-    _CacheSlot,
-)
-
 
 def _sha256(path: Path) -> str:
     return hashlib.sha256(path.read_bytes()).hexdigest()

test/registered/kernels/benchmark/bench_get_block_table.py

@@ -1,7 +1,7 @@
 import torch
 
 from sglang.kernels.jit.benchmark import marker
-from sglang.kernels.jit.minicpm_sala import get_block_table
+from sglang.kernels.ops.minicpm_sala import get_block_table
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/ops/test_get_block_table.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from sglang.kernels.jit.minicpm_sala.get_block_table import get_block_table
+from sglang.kernels.ops.minicpm_sala.get_block_table import get_block_table
 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")

tools/expert_pack/__init__.py

@@ -1 +0,0 @@
-"""DeepSeek expert-pack build, inspection, and validation tools."""