diff --git a/benchmark/bench_attention_sink/bench_attention_sink_triton.py b/benchmark/bench_attention_sink/bench_attention_sink_triton.py index 21bc2a593..f7d563f4d 100644 --- a/benchmark/bench_attention_sink/bench_attention_sink_triton.py +++ b/benchmark/bench_attention_sink/bench_attention_sink_triton.py @@ -3,10 +3,10 @@ import argparse import torch import triton -from sglang.srt.layers.attention.triton_ops.decode_attention import ( +from sglang.kernels.ops.attention.decode_attention import ( decode_attention_fwd_grouped, ) -from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd +from sglang.kernels.ops.attention.extend_attention import extend_attention_fwd # gpt oss head_num = 64 diff --git a/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py b/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py index 94f622bc2..96f839c4f 100644 --- a/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py +++ b/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py @@ -6,8 +6,8 @@ import torch import torch.utils.benchmark as benchmark from flashinfer import BatchDecodeWithPagedKVCacheWrapper +from sglang.kernels.ops.attention.decode_attention import decode_attention_fwd from sglang.srt.layers.attention.flashinfer_backend import should_use_tensor_core -from sglang.srt.layers.attention.triton_ops.decode_attention import decode_attention_fwd def benchmark_forward( diff --git a/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py b/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py index 1c162beca..aa684dfa8 100755 --- a/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py +++ b/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py @@ -58,9 +58,9 @@ from typing import Any, Dict, List, Optional import torch import triton -from sglang.srt.lora.triton_ops.chunked_sgmv_expand import _chunked_lora_expand_kernel -from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel -from sglang.srt.lora.triton_ops.lora_tuning_config import ( +from sglang.kernels.ops.gemm.chunked_sgmv_expand import _chunked_lora_expand_kernel +from sglang.kernels.ops.gemm.chunked_sgmv_shrink import _chunked_lora_shrink_kernel +from sglang.kernels.ops.gemm.lora_tuning_config import ( DEFAULT_EXPAND_CONFIG, DEFAULT_SHRINK_CONFIG, get_lora_config_file_name, @@ -362,9 +362,9 @@ def save_config( "..", "python", "sglang", - "srt", - "lora", - "triton_ops", + "kernels", + "ops", + "gemm", "csgmv_configs", version_dir, ) diff --git a/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py b/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py index 9fd42fb12..340a28ab7 100644 --- a/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py +++ b/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py @@ -5,7 +5,7 @@ import torch.nn.functional as F import triton.testing as tt from sglang.benchmark.bench_utils import run_bench -from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd +from sglang.kernels.ops.attention.extend_attention import extend_attention_fwd def extend_attention_fwd_torch( diff --git a/benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py b/benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py index eb35e8ae2..737308fcc 100644 --- a/benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py +++ b/benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py @@ -19,10 +19,10 @@ import argparse import torch import triton -from sglang.srt.layers.attention.triton_ops.extend_attention import ( +from sglang.kernels.ops.attention.extend_attention import ( extend_attention_fwd, ) -from sglang.srt.layers.attention.triton_ops.verify_splitkv import verify_splitkv_fwd +from sglang.kernels.ops.attention.verify_splitkv import verify_splitkv_fwd from sglang.srt.utils import is_gfx95_supported diff --git a/python/sglang/jit_kernel/dsa/__init__.py b/python/sglang/jit_kernel/dsa/__init__.py index 9c2f92d76..8509eb648 100644 --- a/python/sglang/jit_kernel/dsa/__init__.py +++ b/python/sglang/jit_kernel/dsa/__init__.py @@ -1,13 +1,3 @@ -from sglang.srt.utils import is_cuda - -_is_cuda = is_cuda() - -if _is_cuda: - from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand -else: - CuteDSLPagedMQALogitsRunner = None - pick_dsl_expand = None - from .paged_mqa_logits import ( aiter_paged_mqa_logits, cutedsl_paged_mqa_logits, @@ -15,6 +5,21 @@ from .paged_mqa_logits import ( deepgemm_paged_mqa_logits_split, ) + +def pick_dsl_expand(*args, **kwargs): + from .cutedsl_paged_mqa_logits import pick_dsl_expand as _pick_dsl_expand + + return _pick_dsl_expand(*args, **kwargs) + + +def __getattr__(name: str): + if name == "CuteDSLPagedMQALogitsRunner": + from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner + + return CuteDSLPagedMQALogitsRunner + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "CuteDSLPagedMQALogitsRunner", "pick_dsl_expand", diff --git a/python/sglang/kernels/README.md b/python/sglang/kernels/README.md new file mode 100644 index 000000000..e90f78228 --- /dev/null +++ b/python/sglang/kernels/README.md @@ -0,0 +1,110 @@ +# `sglang.kernels` — unified kernel namespace + +This package is the public in-tree import surface for callable kernels, per +[RFC #29630](https://github.com/sgl-project/sglang/issues/29630). + +```python +from sglang.kernels.ops.layernorm import rmsnorm +from sglang.kernels.ops.activation import silu_and_mul +from sglang.kernels.ops.kvcache import reshape_and_cache_flash +``` + +## Layout + +``` +sglang/kernels/ + spec.py # KernelSpec, KernelBackend, FormatSignature, + # CapabilityRequirement, PlatformInfo + registry.py # process-wide KernelRegistry + register_kernel() + selector.py # heuristic select_kernel() and cached get_kernel() + fused_op.py # BaseFusedOp: per-operator multi-backend contract + ops/ + / # one subpackage per operator group +``` + +Groups populated in this phase: `activation`, `gemm`, `kvcache`, `layernorm`, +`moe`, `quantization`. The remaining groups (`attention`, `communication`, +`diffusion`, `grammar`, `mamba`, `memory`, `sampling`, `spatial`, +`speculative`) are reserved package placeholders whose implementations still +live in `sglang.jit_kernel` / `sgl_kernel` / `triton_ops` and will migrate in +later phases. + +## How it works + +Implementations are not moved yet. Each `ops.` function is a thin +wrapper that forwards to a chosen backend, and every backend is described by a +`KernelSpec` in the registry so alternatives can be inventoried and compared: + +- `register_kernel(KernelSpec(...))` records metadata only — an operator id + (`"."`), a backend, and an import path (`"module:attr"`). No + `torch` or kernel backend is imported, and no JIT compilation is triggered, + until a kernel is actually called. +- `select_kernel(op, backend=None)` resolves an op to its fixed call path. + There is **no** priority ranking or heuristic auto-selection: an op with a + single backend resolves to it; an op with several backends must be resolved + by naming one (`backend=...`). The extra backends are inventory only. +- `get_kernel(op, backend)` resolves and caches the callable; the public + wrappers use it, pinned to the backend whose signature they document. + +The public wrappers currently default to the AOT `sgl_kernel` implementation +(the stable wheel boundary, broadest shape support). The JIT CUDA backend is +registered alongside for inventory; where its signature differs, select it +explicitly, e.g.: + +```python +from sglang.kernels import select_kernel, KernelBackend +jit_rmsnorm = select_kernel("layernorm.rmsnorm", backend=KernelBackend.CUDA_JIT).load() +``` + +## `BaseFusedOp` — the per-operator implementation contract + +Multi-backend operators (currently the `layernorm` and `activation` groups) +are implemented as `BaseFusedOp` subclasses: one logical operator with one +`forward_` method per backend, all sharing one signature behind a +single `forward()`: + +- `forward_native` — **required**; the pure-`torch` correctness reference + every other backend is checked against. +- `forward_torch_compile` — inherited for free as + `torch.compile(forward_native)`. +- `forward_triton` / `forward_cuda_jit` / `forward_cuda_aot` / + `forward_cute_dsl` / `forward_flashinfer` / `forward_deepgemm` — opt-in + overrides. A backend is *available* iff its method is overridden. + +`forward()` auto-selects the best available backend by the class's `priority`, +filtered per call through `backend_eligible()` (a +`CapabilityRequirement`-vs-`PlatformInfo` check, extensible with per-call +shape/dtype gates), and degrades to the native reference when no optimized +backend fits. The public `ops.` functions stay thin wrappers over +module-level instances, so the import surface is unchanged; each instance also +registers all of its backends as `KernelSpec`s so the registry inventory and +`select_kernel(..., backend=...)` keep working. + +What this buys (see the +[RFC discussion](https://github.com/sgl-project/sglang/issues/29630#issuecomment-4920387930)): + +- **Unified correctness testing** — a generic harness enumerates + `available_backends()` and asserts each one matches `forward_native` + (`test/registered/kernels/test_fused_op_gpu_parity.py`); new backends are + picked up automatically. +- **One-switch debugging** — `SGLANG_FORCE_FUSED_OP_BACKEND=torch` (or + `set_fused_op_backend(KernelBackend.TORCH)`) flips *every* fused op to its + reference implementation for numerical-bug bisection. +- **Safe fallbacks** — a missing / ineligible optimized kernel degrades to + `native` instead of scattering `if`/`else` at call sites. +- **Incremental optimization** — land `forward_native` first, add `triton` / + `cuda_jit` / `cuda_aot` later without touching call sites; alternative + implementations of the same op live side by side for A/B. +- **Tracing** — `enable_fused_op_trace()` records every call's op, backend, + and tensor shapes/dtypes, giving an accurate inventory of what a model + actually exercises. + +## Review rule (RFC #29630) + +> SGLang runtime code and tests should import callable kernels from +> `sglang.kernels.ops.*`. + +Implementation work can still happen in `sglang.jit_kernel` or `sgl_kernel`. +When a PR adds a new callable kernel, add a `sglang.kernels.ops.*` entry point +for it, and avoid growing `sglang.jit_kernel` as a long-term public operator +namespace. diff --git a/python/sglang/kernels/__init__.py b/python/sglang/kernels/__init__.py new file mode 100644 index 000000000..8bc34fccf --- /dev/null +++ b/python/sglang/kernels/__init__.py @@ -0,0 +1,72 @@ +"""Unified public kernel namespace for SGLang (RFC #29630). + +SGLang runtime code and tests should import callable kernels from +``sglang.kernels.ops.``, e.g.:: + + from sglang.kernels.ops.layernorm import rmsnorm + from sglang.kernels.ops.activation import silu_and_mul + from sglang.kernels.ops.kvcache import reshape_and_cache_flash + +Implementations still live in ``sglang.jit_kernel`` (JIT CUDA), the +``sgl_kernel`` wheel (AOT CUDA/C++), Triton op modules, etc. The ``ops.*`` +functions are thin wrappers that forward to a chosen backend; the +:data:`~sglang.kernels.registry.registry` provides an inventory of every +backend so alternatives can be compared and selected. Multi-backend operators +are structured as :class:`~sglang.kernels.fused_op.BaseFusedOp` subclasses — +one ``forward_`` method per backend behind a single ``forward()``, +with a required pure-``torch`` ``forward_native`` reference and a +``SGLANG_FORCE_FUSED_OP_BACKEND`` global switch. + +Importing this package (and any ``ops.*`` group) does not import a kernel +backend (``sgl_kernel`` / ``sglang.jit_kernel``) or trigger JIT compilation: +registration is metadata-only and backends are imported lazily on first call. +This keeps the namespace usable for inventory tooling on a CPU-only box. +""" + +from sglang.kernels.fused_op import ( + BaseFusedOp, + FusedOpTraceRecord, + clear_fused_op_trace, + disable_fused_op_trace, + enable_fused_op_trace, + get_fused_op_backend, + get_fused_op_trace, + register_fused_op, + set_fused_op_backend, +) +from sglang.kernels.registry import KernelRegistry, register_kernel, registry +from sglang.kernels.selector import get_kernel, select_kernel +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, + PlatformInfo, +) + +# Importing the operator groups populates the registry (metadata only). Kept +# after the core imports above (and guarded from isort reordering) so those +# modules are fully initialized first. +from sglang.kernels import ops # noqa: E402 # isort: skip + +__all__ = [ + "ops", + "BaseFusedOp", + "CapabilityRequirement", + "FormatSignature", + "FusedOpTraceRecord", + "KernelBackend", + "KernelRegistry", + "KernelSpec", + "PlatformInfo", + "clear_fused_op_trace", + "disable_fused_op_trace", + "enable_fused_op_trace", + "get_fused_op_backend", + "get_fused_op_trace", + "get_kernel", + "register_fused_op", + "register_kernel", + "registry", + "select_kernel", +] diff --git a/python/sglang/kernels/fused_op.py b/python/sglang/kernels/fused_op.py new file mode 100644 index 000000000..e54d7f18d --- /dev/null +++ b/python/sglang/kernels/fused_op.py @@ -0,0 +1,318 @@ +"""Multi-backend operator contract for the unified kernels namespace. + +:class:`BaseFusedOp` is the per-operator implementation object behind the +``sglang.kernels.ops.*`` wrappers (RFC #29630): one logical operator, +implemented once, with multiple interchangeable backends behind a single +``forward()``. + +Each subclass implements one ``forward_`` method per backend it +supports: + +- ``forward_native`` — **required**; the pure-``torch`` correctness reference + every other backend is checked against. +- ``forward_torch_compile`` — provided by the base class as + ``torch.compile(forward_native)``. +- ``forward_triton`` / ``forward_cuda_jit`` / ``forward_cuda_aot`` / + ``forward_cute_dsl`` / ``forward_flashinfer`` / ``forward_deepgemm`` — + opt-in overrides. + +A backend is *available* iff the subclass overrides its method (``native`` and +``torch_compile`` are always available). ``forward()`` picks the best +available backend by :attr:`BaseFusedOp.priority`, filtered per call through +:meth:`BaseFusedOp.backend_eligible` (which checks +:class:`~sglang.kernels.spec.CapabilityRequirement` against the detected +:class:`~sglang.kernels.spec.PlatformInfo`). The ``SGLANG_FORCE_FUSED_OP_BACKEND`` +env var (or :func:`set_fused_op_backend`) forces every fused op onto one +backend — e.g. ``native`` to bisect numerical bugs against the reference +implementations with a single switch. + +Like the rest of ``sglang.kernels``, importing this module (and instantiating +subclasses) never imports a kernel backend (``sgl_kernel`` / +``sglang.jit_kernel``) or triggers JIT compilation; backends are imported +lazily inside the ``forward_`` methods. +""" + +from __future__ import annotations + +import functools +from abc import ABC, abstractmethod +from typing import Any, ClassVar, Dict, List, Mapping, Optional, Tuple + +import msgspec + +from sglang.kernels.registry import register_kernel +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, + PlatformInfo, +) + +# backend -> forward_ method name. +BACKEND_METHODS: Dict[KernelBackend, str] = { + KernelBackend.TORCH: "forward_native", + KernelBackend.TORCH_COMPILE: "forward_torch_compile", + KernelBackend.TRITON: "forward_triton", + KernelBackend.CUDA_JIT: "forward_cuda_jit", + KernelBackend.CUDA_AOT: "forward_cuda_aot", + KernelBackend.CUTE_DSL: "forward_cute_dsl", + KernelBackend.FLASHINFER: "forward_flashinfer", + KernelBackend.DEEPGEMM: "forward_deepgemm", +} + +# best -> fallback. ``torch_compile`` is deliberately absent: auto-selection +# must never trigger a surprise compilation in a serving process; force it +# explicitly when wanted. +DEFAULT_PRIORITY: Tuple[KernelBackend, ...] = ( + KernelBackend.CUDA_AOT, + KernelBackend.CUDA_JIT, + KernelBackend.FLASHINFER, + KernelBackend.DEEPGEMM, + KernelBackend.CUTE_DSL, + KernelBackend.TRITON, + KernelBackend.TORCH, +) + +# Backends every op supports structurally: forward_native is abstract (so a +# concrete subclass always has it) and forward_torch_compile derives from it. +_ALWAYS_AVAILABLE = (KernelBackend.TORCH, KernelBackend.TORCH_COMPILE) + + +@functools.lru_cache(maxsize=1) +def _platform() -> PlatformInfo: + return PlatformInfo.detect() + + +# --- global backend override ------------------------------------------------ + +# Sentinel distinguishing "not resolved yet" from "resolved to None (no force)". +_UNRESOLVED = object() +_forced_backend: Any = _UNRESOLVED + + +def get_fused_op_backend() -> Optional[KernelBackend]: + """The process-wide forced backend, or ``None`` for auto-selection. + + Resolved once from ``SGLANG_FORCE_FUSED_OP_BACKEND`` on first use; tests + and tools flip it afterwards via :func:`set_fused_op_backend`. + """ + global _forced_backend + if _forced_backend is _UNRESOLVED: + from sglang.srt.environ import envs + + value = envs.SGLANG_FORCE_FUSED_OP_BACKEND.get() + _forced_backend = KernelBackend(value) if value is not None else None + return _forced_backend + + +def set_fused_op_backend(backend: Optional[KernelBackend]) -> None: + """Force every :class:`BaseFusedOp` onto ``backend`` (``None`` = auto).""" + global _forced_backend + _forced_backend = backend + + +# --- optional call tracing ---------------------------------------------------- + + +class FusedOpTraceRecord(msgspec.Struct, frozen=True): + """One traced ``forward()`` call: which op ran on which backend, and the + tensor shapes/dtypes it saw.""" + + op: str + backend: str + tensor_args: Tuple[str, ...] # e.g. "torch.bfloat16[128, 4096]" + + +_trace_enabled: bool = False +_trace_records: List[FusedOpTraceRecord] = [] + + +def enable_fused_op_trace() -> None: + """Record every fused-op call (op, backend, tensor shapes/dtypes). + + Gives an accurate inventory of which kernels a model actually exercises + and at what shapes — the raw material for representative-shape test sets. + """ + global _trace_enabled + _trace_enabled = True + + +def disable_fused_op_trace() -> None: + global _trace_enabled + _trace_enabled = False + + +def get_fused_op_trace() -> List[FusedOpTraceRecord]: + return list(_trace_records) + + +def clear_fused_op_trace() -> None: + _trace_records.clear() + + +def _describe_tensors(args: tuple, kwargs: dict) -> Tuple[str, ...]: + import torch + + described = [] + for value in (*args, *kwargs.values()): + if isinstance(value, torch.Tensor): + described.append(f"{value.dtype}[{', '.join(map(str, value.shape))}]") + return tuple(described) + + +# --- the per-operator contract ------------------------------------------------ + + +class BaseFusedOp(ABC): + """One logical operator with interchangeable backends behind ``forward()``. + + Subclasses set :attr:`op` and implement :meth:`forward_native` plus any + optimized ``forward_`` methods. All backend methods of one op + must share the same signature and semantics — each override adapts its + underlying kernel's calling convention so call sites never care which + backend ran. + + Class attributes + ---------------- + op: + Operator id, ``"."`` (e.g. ``"layernorm.rmsnorm"``). + priority: + Backend preference for auto-selection, best first. Defaults to + :data:`DEFAULT_PRIORITY`. + capabilities: + Per-backend :class:`CapabilityRequirement`, consulted by + :meth:`backend_eligible` (and exported into the registry specs). + format_signature: + Data-contract description shared by all backends of this op. + descriptions: + Optional per-backend one-liners for the registry inventory. + """ + + op: ClassVar[str] + priority: ClassVar[Tuple[KernelBackend, ...]] = DEFAULT_PRIORITY + capabilities: ClassVar[Mapping[KernelBackend, CapabilityRequirement]] = {} + format_signature: ClassVar[FormatSignature] = FormatSignature() + descriptions: ClassVar[Mapping[KernelBackend, str]] = {} + + def __init__(self) -> None: + # Cache the structural backend set and the priority-ordered subset once + # so forward() avoids repeated introspection on the hot path. + available = [] + for backend in KernelBackend: + if backend in _ALWAYS_AVAILABLE or self._overrides( + BACKEND_METHODS[backend] + ): + available.append(backend) + self._available: Tuple[KernelBackend, ...] = tuple(available) + self._ordered: Tuple[KernelBackend, ...] = tuple( + b for b in self.priority if b in set(available) + ) + self._compiled_native = None + + def _overrides(self, method_name: str) -> bool: + for klass in type(self).__mro__: + if klass is BaseFusedOp: + return False + if method_name in klass.__dict__: + return True + return False + + # --- backends: native is required; the rest are opt-in overrides --- + + @abstractmethod + def forward_native(self, *args, **kwargs): + """Pure-``torch`` reference implementation (correctness ground truth).""" + + def forward_torch_compile(self, *args, **kwargs): + if self._compiled_native is None: + import torch + + self._compiled_native = torch.compile(self.forward_native) + return self._compiled_native(*args, **kwargs) + + def forward_triton(self, *args, **kwargs): + raise NotImplementedError(f"{self.op}: no triton backend") + + def forward_cuda_jit(self, *args, **kwargs): + raise NotImplementedError(f"{self.op}: no cuda_jit backend") + + def forward_cuda_aot(self, *args, **kwargs): + raise NotImplementedError(f"{self.op}: no cuda_aot backend") + + def forward_cute_dsl(self, *args, **kwargs): + raise NotImplementedError(f"{self.op}: no cute_dsl backend") + + def forward_flashinfer(self, *args, **kwargs): + raise NotImplementedError(f"{self.op}: no flashinfer backend") + + def forward_deepgemm(self, *args, **kwargs): + raise NotImplementedError(f"{self.op}: no deepgemm backend") + + # --- selection --- + + def available_backends(self) -> List[KernelBackend]: + """Backends this op implements (structural check, platform-agnostic).""" + return list(self._available) + + def backend_eligible(self, backend: KernelBackend, *args, **kwargs) -> bool: + """Whether ``backend`` may run *this* call. + + The base implementation checks the backend's + :class:`CapabilityRequirement` against the detected platform. + Subclasses may extend it with per-call shape/dtype gates so + auto-selection bounces to the next backend instead of raising. + """ + capability = self.capabilities.get(backend) + return capability is None or capability.is_satisfied_by(_platform()) + + def _resolve_backend(self, *args, **kwargs) -> KernelBackend: + forced = get_fused_op_backend() + if forced is not None: + return forced + for backend in self._ordered: + if self.backend_eligible(backend, *args, **kwargs): + return backend + return KernelBackend.TORCH + + # --- dispatch --- + + def forward(self, *args, backend: Optional[KernelBackend] = None, **kwargs): + """Run the op on ``backend``, or on the best eligible one when omitted.""" + if backend is None: + backend = self._resolve_backend(*args, **kwargs) + result = getattr(self, BACKEND_METHODS[backend])(*args, **kwargs) + if _trace_enabled: + _trace_records.append( + FusedOpTraceRecord( + op=self.op, + backend=backend.value, + tensor_args=_describe_tensors(args, kwargs), + ) + ) + return result + + __call__ = forward + + +def register_fused_op(instance: BaseFusedOp, module: str, attr: str) -> BaseFusedOp: + """Register every available backend of ``instance`` in the kernel registry. + + ``module``/``attr`` locate the module-level instance so that + ``KernelSpec.load()`` can lazily resolve e.g. + ``":.forward_cuda_aot"`` to the bound backend method. Returns + ``instance`` so group packages can write + ``_RMSNORM = register_fused_op(_RMSNormOp(), __name__, "_RMSNORM")``. + """ + for backend in instance.available_backends(): + register_kernel( + KernelSpec( + op=instance.op, + backend=backend, + target=f"{module}:{attr}.{BACKEND_METHODS[backend]}", + capability=instance.capabilities.get(backend, CapabilityRequirement()), + format_signature=instance.format_signature, + description=instance.descriptions.get(backend, ""), + ) + ) + return instance diff --git a/python/sglang/kernels/ops/__init__.py b/python/sglang/kernels/ops/__init__.py new file mode 100644 index 000000000..5f7979d3c --- /dev/null +++ b/python/sglang/kernels/ops/__init__.py @@ -0,0 +1,40 @@ +"""Public operator groups for the ``sglang.kernels`` namespace. + +Each submodule corresponds to one operator group from RFC #29630. Importing a +group registers its :class:`~sglang.kernels.spec.KernelSpec` metadata and +exposes thin, lazily-dispatched wrapper callables. + +Importing this package eagerly imports every group so the registry is fully +populated for inventory tooling. Group imports are metadata-only and do not +import ``torch`` or a kernel backend. +""" + +from importlib import import_module + +# Operator groups from the RFC's proposed shape. Populated groups expose +# callable wrappers today; the rest are reserved package placeholders that keep +# the namespace shape stable for later phases. +_GROUPS = ( + "activation", + "attention", + "communication", + "diffusion", + "gemm", + "grammar", + "kvcache", + "layernorm", + "mamba", + "memory", + "moe", + "quantization", + "sampling", + "spatial", + "speculative", +) + +for _group in _GROUPS: + import_module(f"{__name__}.{_group}") + +del import_module, _group + +__all__ = list(_GROUPS) diff --git a/python/sglang/kernels/ops/activation/__init__.py b/python/sglang/kernels/ops/activation/__init__.py new file mode 100644 index 000000000..3bd204a5d --- /dev/null +++ b/python/sglang/kernels/ops/activation/__init__.py @@ -0,0 +1,184 @@ +"""Fused gated-activation kernels (``act(x[:h]) * x[h:]``). + +Each operator is a :class:`~sglang.kernels.fused_op.BaseFusedOp` with a +pure-``torch`` reference (``forward_native``) plus AOT (``sgl_kernel``) and +JIT CUDA backends behind one ``(input, out)`` signature. The JIT backend +additionally accepts ``expert_ids`` / ``expert_step`` — call +``forward_cuda_jit`` directly when those are needed. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.kernels.fused_op import BaseFusedOp, register_fused_op +from sglang.kernels.registry import register_kernel +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, +) + +if TYPE_CHECKING: + import torch + +_ACT_DTYPES = ("float16", "bfloat16") +_CUDA = CapabilityRequirement(requires_cuda=True) +_ACT_PRIORITY = ( + KernelBackend.CUDA_AOT, + KernelBackend.CUDA_JIT, + KernelBackend.TORCH, +) + + +class _GatedActivationOp(BaseFusedOp): + """Shared structure for ``act(x[..., :d]) * x[..., d:]`` operators.""" + + # Set by subclasses: sgl_kernel / jit_kernel attr name (same for both). + kernel_attr: str + + priority = _ACT_PRIORITY + capabilities = { + KernelBackend.CUDA_AOT: _CUDA, + KernelBackend.CUDA_JIT: _CUDA, + } + format_signature = FormatSignature( + supported_dtypes=_ACT_DTYPES, + description="gated activation; returns tensor", + ) + + def _act(self, gate: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + def forward_native( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None + ) -> torch.Tensor: + d = input.shape[-1] // 2 + result = self._act(input[..., :d]) * input[..., d:] + if out is None: + return result + out.copy_(result) + return out + + def forward_cuda_aot( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None + ) -> torch.Tensor: + import sgl_kernel + + return getattr(sgl_kernel, self.kernel_attr)(input, out) + + def forward_cuda_jit( + self, + input: torch.Tensor, + out: Optional[torch.Tensor] = None, + expert_ids: Optional[torch.Tensor] = None, + expert_step: int = 1, + ) -> torch.Tensor: + import sglang.jit_kernel.activation as jit_activation + + return getattr(jit_activation, self.kernel_attr)( + input, out, expert_ids, expert_step + ) + + +class SiluAndMulOp(_GatedActivationOp): + """``out = silu(input[..., :d]) * input[..., d:]`` with ``d = input.shape[-1] // 2``.""" + + op = "activation.silu_and_mul" + kernel_attr = "silu_and_mul" + descriptions = { + KernelBackend.CUDA_AOT: "silu_and_mul (sgl_kernel wheel).", + KernelBackend.CUDA_JIT: "silu_and_mul (sglang.jit_kernel).", + KernelBackend.TORCH: "silu_and_mul (pure-torch reference).", + } + + def _act(self, gate: torch.Tensor) -> torch.Tensor: + import torch.nn.functional as F + + return F.silu(gate) + + +class GeluAndMulOp(_GatedActivationOp): + """``out = gelu(input[..., :d]) * input[..., d:]`` (erf-based GELU).""" + + op = "activation.gelu_and_mul" + kernel_attr = "gelu_and_mul" + descriptions = { + KernelBackend.CUDA_AOT: "gelu_and_mul (sgl_kernel wheel).", + KernelBackend.CUDA_JIT: "gelu_and_mul (sglang.jit_kernel).", + KernelBackend.TORCH: "gelu_and_mul (pure-torch reference).", + } + + def _act(self, gate: torch.Tensor) -> torch.Tensor: + import torch.nn.functional as F + + return F.gelu(gate, approximate="none") + + +class GeluTanhAndMulOp(_GatedActivationOp): + """``out = gelu_tanh(input[..., :d]) * input[..., d:]`` (tanh-approximated GELU).""" + + op = "activation.gelu_tanh_and_mul" + kernel_attr = "gelu_tanh_and_mul" + descriptions = { + KernelBackend.CUDA_AOT: "gelu_tanh_and_mul (sgl_kernel wheel).", + KernelBackend.CUDA_JIT: "gelu_tanh_and_mul (sglang.jit_kernel).", + KernelBackend.TORCH: "gelu_tanh_and_mul (pure-torch reference).", + } + + def _act(self, gate: torch.Tensor) -> torch.Tensor: + import torch.nn.functional as F + + return F.gelu(gate, approximate="tanh") + + +_SILU_AND_MUL = register_fused_op(SiluAndMulOp(), __name__, "_SILU_AND_MUL") +_GELU_AND_MUL = register_fused_op(GeluAndMulOp(), __name__, "_GELU_AND_MUL") +_GELU_TANH_AND_MUL = register_fused_op( + GeluTanhAndMulOp(), __name__, "_GELU_TANH_AND_MUL" +) + + +def silu_and_mul( + input: torch.Tensor, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + """``out = silu(input[..., :d]) * input[..., d:]`` with ``d = input.shape[-1] // 2``.""" + return _SILU_AND_MUL(input, out) + + +def gelu_and_mul( + input: torch.Tensor, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + """``out = gelu(input[..., :d]) * input[..., d:]``.""" + return _GELU_AND_MUL(input, out) + + +def gelu_tanh_and_mul( + input: torch.Tensor, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + """``out = gelu_tanh(input[..., :d]) * input[..., d:]``.""" + return _GELU_TANH_AND_MUL(input, out) + + +__all__ = [ + "SiluAndMulOp", + "GeluAndMulOp", + "GeluTanhAndMulOp", + "silu_and_mul", + "gelu_and_mul", + "gelu_tanh_and_mul", +] + + +# Triton kernel migrated into this group (from layers/triton_ops/softcap); +# registered for inventory. Import it from its module. +for _fn in ("softcap_out", "softcap_inplace_logits"): + register_kernel( + KernelSpec( + op=f"activation.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.activation.softcap:{_fn}", + ) + ) +del _fn diff --git a/python/sglang/srt/layers/triton_ops/softcap.py b/python/sglang/kernels/ops/activation/softcap.py similarity index 100% rename from python/sglang/srt/layers/triton_ops/softcap.py rename to python/sglang/kernels/ops/activation/softcap.py diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py new file mode 100644 index 000000000..1410c3070 --- /dev/null +++ b/python/sglang/kernels/ops/attention/__init__.py @@ -0,0 +1,43 @@ +"""Attention compute kernels (Triton): decode / extend / prefill / metadata. + +The Triton kernels migrated here live in this package +(``sglang.kernels.ops.attention.``); import them from there. Their +``KernelSpec`` metadata is registered below for inventory (backend = Triton). +KV-cache index/write kernels went to the ``kvcache`` group instead. +""" + +from sglang.kernels.registry import register_kernel +from sglang.kernels.spec import KernelBackend, KernelSpec + +# (module, public_fn) migrated from layers/attention/triton_ops + model_executor. +_TRITON_KERNELS = [ + ("decode_attention", "decode_attention_fwd"), + ("extend_attention", "extend_attention_fwd"), + ("extend_attention", "build_unified_kv_indices"), + ("prefill_attention", "context_attention_fwd"), + ("merge_state", "merge_state_triton"), + ("metadata", "get_num_kv_splits_triton"), + ("metadata", "prepare_swa_spec_page_table_triton"), + ("metadata", "normal_decode_set_metadata"), + ("dsa_metadata", "fused_dsa_decode_metadata"), + ("dsa_metadata", "fused_dsa_target_verify_metadata"), + ("dsa_metadata", "fused_dsa_draft_extend_metadata"), + ("rocm_mla_decode_rope", "decode_attention_fwd_grouped_rope"), + ("verify_splitkv", "verify_splitkv_fwd"), + ("pad", "pad_sequence_with_mask"), + ("pad", "pad_draft_extend_query"), + ("pad", "unpad_draft_extend_output"), + ("pad", "seqlens_expand_triton"), + ("position", "compute_position_triton"), +] +for _mod, _fn in _TRITON_KERNELS: + register_kernel( + KernelSpec( + op=f"attention.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.attention.{_mod}:{_fn}", + ) + ) +del _mod, _fn + +__all__ = [] diff --git a/python/sglang/srt/layers/attention/triton_ops/decode_attention.py b/python/sglang/kernels/ops/attention/decode_attention.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/decode_attention.py rename to python/sglang/kernels/ops/attention/decode_attention.py diff --git a/python/sglang/srt/layers/attention/triton_ops/dsa_metadata.py b/python/sglang/kernels/ops/attention/dsa_metadata.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/dsa_metadata.py rename to python/sglang/kernels/ops/attention/dsa_metadata.py diff --git a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py b/python/sglang/kernels/ops/attention/extend_attention.py similarity index 99% rename from python/sglang/srt/layers/attention/triton_ops/extend_attention.py rename to python/sglang/kernels/ops/attention/extend_attention.py index 37d96d3ac..97f94db55 100644 --- a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py +++ b/python/sglang/kernels/ops/attention/extend_attention.py @@ -20,8 +20,8 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.triton_ops.decode_attention import _extract_kv_strides -from sglang.srt.layers.attention.triton_ops.prefill_attention import ( +from sglang.kernels.ops.attention.decode_attention import _extract_kv_strides +from sglang.kernels.ops.attention.prefill_attention import ( context_attention_fwd, ) from sglang.srt.utils import is_cuda, is_gfx95_supported, is_hip diff --git a/python/sglang/srt/layers/attention/triton_ops/merge_state.py b/python/sglang/kernels/ops/attention/merge_state.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/merge_state.py rename to python/sglang/kernels/ops/attention/merge_state.py diff --git a/python/sglang/srt/layers/attention/triton_ops/metadata.py b/python/sglang/kernels/ops/attention/metadata.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/metadata.py rename to python/sglang/kernels/ops/attention/metadata.py diff --git a/python/sglang/srt/layers/attention/triton_ops/pad.py b/python/sglang/kernels/ops/attention/pad.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/pad.py rename to python/sglang/kernels/ops/attention/pad.py diff --git a/python/sglang/srt/model_executor/triton_ops/position.py b/python/sglang/kernels/ops/attention/position.py similarity index 100% rename from python/sglang/srt/model_executor/triton_ops/position.py rename to python/sglang/kernels/ops/attention/position.py diff --git a/python/sglang/srt/layers/attention/triton_ops/prefill_attention.py b/python/sglang/kernels/ops/attention/prefill_attention.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/prefill_attention.py rename to python/sglang/kernels/ops/attention/prefill_attention.py diff --git a/python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py b/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py similarity index 99% rename from python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py rename to python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py index 8cd397e30..837dacee5 100644 --- a/python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py +++ b/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py @@ -23,7 +23,7 @@ It supports page size = 1. import triton import triton.language as tl -from sglang.srt.layers.attention.triton_ops.decode_attention import ( +from sglang.kernels.ops.attention.decode_attention import ( _decode_softmax_reducev_fwd, ) diff --git a/python/sglang/srt/layers/attention/triton_ops/verify_splitkv.py b/python/sglang/kernels/ops/attention/verify_splitkv.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/verify_splitkv.py rename to python/sglang/kernels/ops/attention/verify_splitkv.py diff --git a/python/sglang/kernels/ops/communication/__init__.py b/python/sglang/kernels/ops/communication/__init__.py new file mode 100644 index 000000000..1d20cb8ff --- /dev/null +++ b/python/sglang/kernels/ops/communication/__init__.py @@ -0,0 +1,12 @@ +"""Collective-communication kernels (custom all-reduce, ...). + +Reserved group in the ``sglang.kernels`` namespace (RFC #29630). No thin +wrappers are exposed here: the collective ops (custom all-reduce and friends) +are stateful — they manage workspaces / IPC handles and are driven through a +``CustomAllreduce``-style object and ``torch.ops.sgl_kernel.*`` bindings rather +than standalone callable kernels, so a thin ``sglang.kernels.ops`` forwarder +would be misleading. Import them from ``sgl_kernel`` / ``sglang.jit_kernel`` +directly until a proper stateful-op interface is designed. +""" + +__all__ = [] diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py new file mode 100644 index 000000000..12c2d42bc --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -0,0 +1,109 @@ +"""Diffusion-model kernels (group-norm+silu, residual-gate-add, qk-norm+rope). + +These are JIT CUDA kernels; the wrappers forward to ``sglang.jit_kernel.diffusion``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, +) + +if TYPE_CHECKING: + import torch + from torch import nn + +_CUDA = CapabilityRequirement(requires_cuda=True) + +register_kernel( + KernelSpec( + op="diffusion.apply_group_norm_silu", + backend=KernelBackend.CUDA_JIT, + target="sglang.jit_kernel.diffusion.group_norm_silu:apply_group_norm_silu", + capability=_CUDA, + format_signature=FormatSignature(description="fused GroupNorm + SiLU"), + description="Fused group-norm + SiLU (sglang.jit_kernel).", + ) +) +register_kernel( + KernelSpec( + op="diffusion.residual_gate_add", + backend=KernelBackend.CUDA_JIT, + target="sglang.jit_kernel.diffusion.residual_gate_add:residual_gate_add_cuda", + capability=_CUDA, + format_signature=FormatSignature(description="residual + gate * update"), + description="Fused residual gate-add (sglang.jit_kernel).", + ) +) +register_kernel( + KernelSpec( + op="diffusion.fused_inplace_qknorm_rope", + backend=KernelBackend.CUDA_JIT, + target="sglang.jit_kernel.diffusion.qknorm_rope:fused_inplace_qknorm_rope", + capability=_CUDA, + format_signature=FormatSignature( + in_place=True, description="fused in-place QK-norm + RoPE" + ), + description="Fused QK-norm + RoPE (sglang.jit_kernel).", + ) +) + + +def apply_group_norm_silu( + x: torch.Tensor, norm: nn.Module, activation: nn.Module +) -> torch.Tensor: + """Fused GroupNorm + SiLU (falls back to eager when unsupported).""" + return get_kernel("diffusion.apply_group_norm_silu", KernelBackend.CUDA_JIT)( + x, norm, activation + ) + + +def residual_gate_add( + residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor +) -> torch.Tensor: + """Fused ``residual + gate * update``.""" + return get_kernel("diffusion.residual_gate_add", KernelBackend.CUDA_JIT)( + residual, update, gate + ) + + +def fused_inplace_qknorm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + *, + is_neox: bool, + eps: float = 1e-6, + head_dim: int = 0, + rope_dim: int = 0, +) -> None: + """Fused in-place QK RMS-norm + RoPE.""" + return get_kernel("diffusion.fused_inplace_qknorm_rope", KernelBackend.CUDA_JIT)( + q, + k, + q_weight, + k_weight, + cos_sin_cache, + positions, + is_neox=is_neox, + eps=eps, + head_dim=head_dim, + rope_dim=rope_dim, + ) + + +__all__ = [ + "apply_group_norm_silu", + "residual_gate_add", + "fused_inplace_qknorm_rope", +] diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py new file mode 100644 index 000000000..d11f7ed0d --- /dev/null +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -0,0 +1,138 @@ +"""GEMM and fused-GEMM kernels.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, +) + +if TYPE_CHECKING: + import torch + +_CUDA = CapabilityRequirement(requires_cuda=True) + +register_kernel( + KernelSpec( + op="gemm.fp8_scaled_mm", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel:fp8_scaled_mm", + format_signature=FormatSignature( + supported_dtypes=("float8_e4m3fn",), + description="C = (A_fp8 @ B_fp8) * scales_a * scales_b (+ bias)", + ), + description="FP8 scaled matmul (sgl_kernel wheel).", + ) +) +register_kernel( + KernelSpec( + op="gemm.dsv3_fused_a_gemm", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel:dsv3_fused_a_gemm", + format_signature=FormatSignature( + supported_dtypes=("bfloat16",), + description="DeepSeek-V3 fused QKV-A GEMM", + ), + description="DeepSeek-V3 fused-A GEMM (sgl_kernel wheel).", + ) +) +register_kernel( + KernelSpec( + op="gemm.dsv3_fused_a_gemm", + backend=KernelBackend.CUDA_JIT, + target="sglang.jit_kernel.dsv3_fused_a_gemm:dsv3_fused_a_gemm", + capability=_CUDA, + format_signature=FormatSignature( + supported_dtypes=("bfloat16",), + description="DeepSeek-V3 fused QKV-A GEMM (drop-in with AOT signature)", + ), + description="DeepSeek-V3 fused-A GEMM (sglang.jit_kernel).", + ) +) +register_kernel( + KernelSpec( + op="gemm.dsv3_router_gemm", + backend=KernelBackend.CUDA_JIT, + target="sglang.jit_kernel.dsv3_router_gemm:dsv3_router_gemm", + capability=_CUDA, + format_signature=FormatSignature( + supported_dtypes=("bfloat16",), + description="DeepSeek-V3 router GEMM; num_tokens in [1, 16]", + ), + description="DeepSeek-V3 router GEMM (sglang.jit_kernel, JIT-only).", + ) +) + + +def fp8_scaled_mm( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + scales_a: torch.Tensor, + scales_b: torch.Tensor, + out_dtype: torch.dtype, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """FP8 scaled matmul: ``(mat_a @ mat_b) * scales_a * scales_b (+ bias)``.""" + return get_kernel("gemm.fp8_scaled_mm", KernelBackend.CUDA_AOT)( + mat_a, mat_b, scales_a, scales_b, out_dtype, bias + ) + + +def dsv3_fused_a_gemm( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + output: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """DeepSeek-V3 fused QKV-A GEMM.""" + return get_kernel("gemm.dsv3_fused_a_gemm", KernelBackend.CUDA_AOT)( + mat_a, mat_b, output + ) + + +def dsv3_router_gemm( + hidden_states: torch.Tensor, + router_weights: torch.Tensor, + out_dtype: Optional[torch.dtype] = None, + output: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """DeepSeek-V3 router GEMM (JIT-backed). ``out_dtype`` defaults to bfloat16.""" + impl = get_kernel("gemm.dsv3_router_gemm", KernelBackend.CUDA_JIT) + if out_dtype is None: + return impl(hidden_states, router_weights, output=output) + return impl(hidden_states, router_weights, out_dtype, output) + + +__all__ = ["fp8_scaled_mm", "dsv3_fused_a_gemm", "dsv3_router_gemm"] + + +# LoRA SGMV Triton kernels migrated into this group (from lora/triton_ops); +# registered for inventory. Import them from their modules. +_TRITON_KERNELS = [ + ("chunked_embedding_lora_a", "chunked_embedding_lora_a_forward"), + ("chunked_sgmv_expand", "chunked_sgmv_lora_expand_forward"), + ("chunked_sgmv_shrink", "chunked_sgmv_lora_shrink_forward"), + ("embedding_lora_a", "embedding_lora_a_fwd"), + ("gate_up_lora_b", "gate_up_lora_b_fwd"), + ("qkv_lora_b", "qkv_lora_b_fwd"), + ("sgemm_lora_a", "sgemm_lora_a_fwd"), + ("sgemm_lora_b", "sgemm_lora_b_fwd"), + ("kv_b_lora_absorbed", "step_a_q_fwd"), + ("kv_b_lora_absorbed", "step_b_q_fwd"), + ("kv_b_lora_absorbed", "step_a_v_fwd"), + ("kv_b_lora_absorbed", "step_b_v_fwd"), +] +for _mod, _fn in _TRITON_KERNELS: + register_kernel( + KernelSpec( + op=f"gemm.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.gemm.{_mod}:{_fn}", + ) + ) +del _mod, _fn diff --git a/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py b/python/sglang/kernels/ops/gemm/chunked_embedding_lora_a.py similarity index 100% rename from python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py rename to python/sglang/kernels/ops/gemm/chunked_embedding_lora_a.py diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py b/python/sglang/kernels/ops/gemm/chunked_sgmv_expand.py similarity index 98% rename from python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py rename to python/sglang/kernels/ops/gemm/chunked_sgmv_expand.py index e1968ff99..673d340c6 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py +++ b/python/sglang/kernels/ops/gemm/chunked_sgmv_expand.py @@ -4,7 +4,7 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.triton_ops.lora_tuning_config import get_lora_expand_config +from sglang.kernels.ops.gemm.lora_tuning_config import get_lora_expand_config from sglang.srt.lora.utils import LoRABatchInfo from sglang.srt.utils import cached_triton_kernel diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py b/python/sglang/kernels/ops/gemm/chunked_sgmv_shrink.py similarity index 98% rename from python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py rename to python/sglang/kernels/ops/gemm/chunked_sgmv_shrink.py index 8788a65de..9d7a822bf 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py +++ b/python/sglang/kernels/ops/gemm/chunked_sgmv_shrink.py @@ -2,7 +2,7 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.triton_ops.lora_tuning_config import get_lora_shrink_config +from sglang.kernels.ops.gemm.lora_tuning_config import get_lora_shrink_config from sglang.srt.lora.utils import LoRABatchInfo from sglang.srt.utils import cached_triton_kernel diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json b/python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json similarity index 100% rename from python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json rename to python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json b/python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json similarity index 100% rename from python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json rename to python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json b/python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json similarity index 100% rename from python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json rename to python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json b/python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json similarity index 100% rename from python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json rename to python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json b/python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json similarity index 100% rename from python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json rename to python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json b/python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json similarity index 100% rename from python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json rename to python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json b/python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json similarity index 100% rename from python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json rename to python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json diff --git a/python/sglang/srt/lora/triton_ops/embedding_lora_a.py b/python/sglang/kernels/ops/gemm/embedding_lora_a.py similarity index 100% rename from python/sglang/srt/lora/triton_ops/embedding_lora_a.py rename to python/sglang/kernels/ops/gemm/embedding_lora_a.py diff --git a/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py b/python/sglang/kernels/ops/gemm/gate_up_lora_b.py similarity index 98% rename from python/sglang/srt/lora/triton_ops/gate_up_lora_b.py rename to python/sglang/kernels/ops/gemm/gate_up_lora_b.py index 16ade8b44..bc89fac40 100644 --- a/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py +++ b/python/sglang/kernels/ops/gemm/gate_up_lora_b.py @@ -2,7 +2,7 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions +from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo diff --git a/python/sglang/srt/lora/triton_ops/kernel_utils.py b/python/sglang/kernels/ops/gemm/kernel_utils.py similarity index 100% rename from python/sglang/srt/lora/triton_ops/kernel_utils.py rename to python/sglang/kernels/ops/gemm/kernel_utils.py diff --git a/python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py b/python/sglang/kernels/ops/gemm/kv_b_lora_absorbed.py similarity index 99% rename from python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py rename to python/sglang/kernels/ops/gemm/kv_b_lora_absorbed.py index e4ed6a35b..99a771081 100644 --- a/python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py +++ b/python/sglang/kernels/ops/gemm/kv_b_lora_absorbed.py @@ -48,7 +48,7 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions +from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo # --------------------------------------------------------------------------- diff --git a/python/sglang/srt/lora/triton_ops/lora_tuning_config.py b/python/sglang/kernels/ops/gemm/lora_tuning_config.py similarity index 99% rename from python/sglang/srt/lora/triton_ops/lora_tuning_config.py rename to python/sglang/kernels/ops/gemm/lora_tuning_config.py index 33e9e72ed..d7335694f 100644 --- a/python/sglang/srt/lora/triton_ops/lora_tuning_config.py +++ b/python/sglang/kernels/ops/gemm/lora_tuning_config.py @@ -20,7 +20,7 @@ Usage: python3 benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \ --model Qwen/Qwen3-Embedding-0.6B --max-lora-rank 64 - # Configs saved to python/sglang/srt/lora/triton_ops/configs/ + # Configs saved to python/sglang/kernels/ops/gemm/configs/ # Server automatically picks them up: python3 -m sglang.launch_server --model ... --enable-lora --lora-backend csgmv diff --git a/python/sglang/srt/lora/triton_ops/qkv_lora_b.py b/python/sglang/kernels/ops/gemm/qkv_lora_b.py similarity index 98% rename from python/sglang/srt/lora/triton_ops/qkv_lora_b.py rename to python/sglang/kernels/ops/gemm/qkv_lora_b.py index d43f0c64a..1b1cff2b9 100644 --- a/python/sglang/srt/lora/triton_ops/qkv_lora_b.py +++ b/python/sglang/kernels/ops/gemm/qkv_lora_b.py @@ -2,7 +2,7 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions +from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo diff --git a/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py b/python/sglang/kernels/ops/gemm/sgemm_lora_a.py similarity index 98% rename from python/sglang/srt/lora/triton_ops/sgemm_lora_a.py rename to python/sglang/kernels/ops/gemm/sgemm_lora_a.py index 0dd3e5bbb..549620e42 100644 --- a/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py +++ b/python/sglang/kernels/ops/gemm/sgemm_lora_a.py @@ -2,7 +2,7 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions +from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo diff --git a/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py b/python/sglang/kernels/ops/gemm/sgemm_lora_b.py similarity index 98% rename from python/sglang/srt/lora/triton_ops/sgemm_lora_b.py rename to python/sglang/kernels/ops/gemm/sgemm_lora_b.py index fc7f844e2..c1cfb2ea4 100644 --- a/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py +++ b/python/sglang/kernels/ops/gemm/sgemm_lora_b.py @@ -2,7 +2,7 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions +from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo diff --git a/python/sglang/kernels/ops/gemm/trtllm_lora_temp/__init__.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/__init__.py new file mode 100644 index 000000000..fdc403e21 --- /dev/null +++ b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/__init__.py @@ -0,0 +1,3 @@ +"""Experimental TRT-LLM LoRA kernel variants (gated by ``SGLANG_EXPERIMENTAL_LORA_OPTI`` / ``lora_envs``). + +Migrated from ``sglang.srt.lora.trtllm_lora_temp.triton_ops`` (RFC #29630).""" diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/gate_up_lora_b.py similarity index 99% rename from python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py rename to python/sglang/kernels/ops/gemm/trtllm_lora_temp/gate_up_lora_b.py index b8a76de26..a2893c2c0 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py +++ b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/gate_up_lora_b.py @@ -2,11 +2,11 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs -from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import ( +from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import ( _resolve_token_positions, get_pdl_launch_metadata, ) +from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs from sglang.srt.lora.utils import LoRABatchInfo # Minimum total_tokens * rank for the single-adapter cuBLAS path; below this diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kernel_utils.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py similarity index 100% rename from python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kernel_utils.py rename to python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/kv_b_lora_absorbed.py similarity index 99% rename from python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py rename to python/sglang/kernels/ops/gemm/trtllm_lora_temp/kv_b_lora_absorbed.py index a69d7b189..4ee5ce7f4 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py +++ b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/kv_b_lora_absorbed.py @@ -48,11 +48,11 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs -from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import ( +from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import ( _resolve_token_positions, get_pdl_launch_metadata, ) +from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs from sglang.srt.lora.utils import LoRABatchInfo # --------------------------------------------------------------------------- diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/qkv_lora_b.py similarity index 99% rename from python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py rename to python/sglang/kernels/ops/gemm/trtllm_lora_temp/qkv_lora_b.py index 628d1fbb0..4c8e31cb9 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py +++ b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/qkv_lora_b.py @@ -4,11 +4,11 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs -from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import ( +from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import ( _resolve_token_positions, get_pdl_launch_metadata, ) +from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs from sglang.srt.lora.utils import LoRABatchInfo # Minimum max_len (longest segment) for the single-adapter cuBLAS path; below diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_a.py similarity index 99% rename from python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py rename to python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_a.py index c109fcf3f..293f3ae6c 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py +++ b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_a.py @@ -4,11 +4,11 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs -from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import ( +from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import ( _resolve_token_positions, get_pdl_launch_metadata, ) +from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs from sglang.srt.lora.utils import LoRABatchInfo diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_b.py similarity index 97% rename from python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py rename to python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_b.py index abcc49d96..a80b4c78b 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py +++ b/python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_b.py @@ -2,14 +2,14 @@ import torch import triton import triton.language as tl -from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs -from sglang.srt.lora.trtllm_lora_temp.triton_ops.gate_up_lora_b import ( +from sglang.kernels.ops.gemm.trtllm_lora_temp.gate_up_lora_b import ( _CUBLAS_MIN_S_RANK, ) -from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import ( +from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import ( _resolve_token_positions, get_pdl_launch_metadata, ) +from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs from sglang.srt.lora.utils import LoRABatchInfo diff --git a/python/sglang/kernels/ops/grammar/__init__.py b/python/sglang/kernels/ops/grammar/__init__.py new file mode 100644 index 000000000..183c57557 --- /dev/null +++ b/python/sglang/kernels/ops/grammar/__init__.py @@ -0,0 +1,26 @@ +"""Constrained-decoding / grammar kernels (Triton). + +The Triton kernels migrated here live in this package +(``sglang.kernels.ops.grammar.``); import them from there. Their +``KernelSpec`` metadata is registered below for inventory (backend = Triton). +""" + +from sglang.kernels.registry import register_kernel +from sglang.kernels.spec import KernelBackend, KernelSpec + +# (module, public_fn) migrated from constrained/triton_ops. +_TRITON_KERNELS = [ + ("bitmask_ops", "apply_token_bitmask_inplace_triton"), + ("token_filter_ops", "set_token_filter_triton"), +] +for _mod, _fn in _TRITON_KERNELS: + register_kernel( + KernelSpec( + op=f"grammar.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.grammar.{_mod}:{_fn}", + ) + ) +del _mod, _fn + +__all__ = [] diff --git a/python/sglang/srt/constrained/triton_ops/bitmask_ops.py b/python/sglang/kernels/ops/grammar/bitmask_ops.py similarity index 100% rename from python/sglang/srt/constrained/triton_ops/bitmask_ops.py rename to python/sglang/kernels/ops/grammar/bitmask_ops.py diff --git a/python/sglang/srt/constrained/triton_ops/token_filter_ops.py b/python/sglang/kernels/ops/grammar/token_filter_ops.py similarity index 100% rename from python/sglang/srt/constrained/triton_ops/token_filter_ops.py rename to python/sglang/kernels/ops/grammar/token_filter_ops.py diff --git a/python/sglang/kernels/ops/kvcache/__init__.py b/python/sglang/kernels/ops/kvcache/__init__.py new file mode 100644 index 000000000..b6501baaa --- /dev/null +++ b/python/sglang/kernels/ops/kvcache/__init__.py @@ -0,0 +1,89 @@ +"""KV-cache write/transfer kernels. + +This group wraps the Triton ``reshape_and_cache`` launcher, whose implementation +now lives in this package (``sglang.kernels.ops.kvcache.cache_ops``) after being +migrated out of ``sglang.srt.layers.attention.triton_ops`` (RFC #29630). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec + +if TYPE_CHECKING: + import torch + +register_kernel( + KernelSpec( + op="kvcache.reshape_and_cache_flash", + backend=KernelBackend.TRITON, + target="sglang.kernels.ops.kvcache.cache_ops:launch_reshape_and_cache_flash", + format_signature=FormatSignature( + in_place=True, + description="write token-major K/V into paged KV cache layout", + ), + description="Reshape-and-cache (Triton launcher).", + ) +) + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + swa_slot_mapping: Optional[torch.Tensor] = None, + k_scale: Optional[torch.Tensor] = None, + v_scale: Optional[torch.Tensor] = None, +) -> None: + """Write token-major ``key``/``value`` into paged KV cache layout.""" + return get_kernel("kvcache.reshape_and_cache_flash", KernelBackend.TRITON)( + key, + value, + key_cache, + value_cache, + slot_mapping, + swa_slot_mapping, + k_scale, + v_scale, + ) + + +__all__ = ["reshape_and_cache_flash"] + + +# Other Triton kernels migrated into this group (from attention/mem_cache +# triton_ops); registered for inventory. Import them from their modules. +_TRITON_KERNELS = [ + ("cache_ops", "concat_and_cast_mha_k_triton"), + ("cache_ops", "launch_reshape_and_cache_flash"), + ("kv_indices", "create_flashinfer_kv_indices_triton"), + ("kv_indices", "create_flashmla_kv_indices_triton"), + ("kv_indices", "create_chunked_prefix_cache_kv_indices"), + ("kv_indices", "get_num_kv_index_blocks_flashmla"), + ("kv_indices", "get_num_page_per_block_flashmla"), + ("rope_cache", "fused_qk_rope_reshape_and_cache"), + ("trtllm_fp8_kv_kernel", "fused_fp8_set_kv_buffer"), + ("trtllm_mha_page_table", "build_trtllm_mha_page_table"), + ("trtllm_mha_graph_metadata", "update_trtllm_mha_graph_metadata"), + ("aiter_unified_attention", "scatter_ragged_to_page_table_kernel"), + ("aiter_unified_attention", "scatter_req_to_token_to_page_table_kernel"), + ("cache_move", "store_cache_4d"), + ("cache_move", "set_kv_buffer_prefix_valid_tiled"), + ("cache_move", "copy_all_layer_kv_cache_tiled"), + ("mla_buffer", "set_mla_kv_buffer_triton"), + ("mla_buffer", "get_mla_kv_buffer_triton"), +] +for _mod, _fn in _TRITON_KERNELS: + register_kernel( + KernelSpec( + op=f"kvcache.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.kvcache.{_mod}:{_fn}", + ) + ) +del _mod, _fn diff --git a/python/sglang/srt/layers/attention/triton_ops/aiter_unified_attention.py b/python/sglang/kernels/ops/kvcache/aiter_unified_attention.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/aiter_unified_attention.py rename to python/sglang/kernels/ops/kvcache/aiter_unified_attention.py diff --git a/python/sglang/srt/mem_cache/triton_ops/cache_move.py b/python/sglang/kernels/ops/kvcache/cache_move.py similarity index 100% rename from python/sglang/srt/mem_cache/triton_ops/cache_move.py rename to python/sglang/kernels/ops/kvcache/cache_move.py diff --git a/python/sglang/srt/layers/attention/triton_ops/cache_ops.py b/python/sglang/kernels/ops/kvcache/cache_ops.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/cache_ops.py rename to python/sglang/kernels/ops/kvcache/cache_ops.py diff --git a/python/sglang/srt/layers/attention/triton_ops/kv_indices.py b/python/sglang/kernels/ops/kvcache/kv_indices.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/kv_indices.py rename to python/sglang/kernels/ops/kvcache/kv_indices.py diff --git a/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py b/python/sglang/kernels/ops/kvcache/mla_buffer.py similarity index 100% rename from python/sglang/srt/mem_cache/triton_ops/mla_buffer.py rename to python/sglang/kernels/ops/kvcache/mla_buffer.py diff --git a/python/sglang/srt/layers/attention/triton_ops/rope_cache.py b/python/sglang/kernels/ops/kvcache/rope_cache.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/rope_cache.py rename to python/sglang/kernels/ops/kvcache/rope_cache.py diff --git a/python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py b/python/sglang/kernels/ops/kvcache/trtllm_fp8_kv_kernel.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py rename to python/sglang/kernels/ops/kvcache/trtllm_fp8_kv_kernel.py diff --git a/python/sglang/srt/layers/attention/triton_ops/trtllm_mha_graph_metadata.py b/python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/trtllm_mha_graph_metadata.py rename to python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py diff --git a/python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py b/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py similarity index 100% rename from python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py rename to python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py new file mode 100644 index 000000000..ae22499ea --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -0,0 +1,338 @@ +"""Layer-normalization kernels. + +Each operator is a :class:`~sglang.kernels.fused_op.BaseFusedOp` with a +pure-``torch`` reference (``forward_native``) plus optimized CUDA backends, +all behind one signature. The public module-level functions are thin wrappers +over module-level instances; auto-selection prefers the AOT ``sgl_kernel`` +implementation on CUDA and falls back to the native reference elsewhere. +Pick a specific backend with e.g. +``_RMSNORM.forward(x, w, backend=KernelBackend.CUDA_JIT)`` or globally via +``SGLANG_FORCE_FUSED_OP_BACKEND``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.kernels.fused_op import BaseFusedOp, register_fused_op +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, +) + +if TYPE_CHECKING: + import torch + +_NORM_DTYPES = ("float16", "bfloat16") +_CUDA = CapabilityRequirement(requires_cuda=True) +_NORM_PRIORITY = ( + KernelBackend.CUDA_AOT, + KernelBackend.CUDA_JIT, + KernelBackend.TORCH, +) + + +class RMSNormOp(BaseFusedOp): + """``out = (input / RMS(input)) * weight``; returns a tensor. + + ``enable_pdl`` is honored by the AOT backend only. + """ + + op = "layernorm.rmsnorm" + priority = _NORM_PRIORITY + capabilities = { + KernelBackend.CUDA_AOT: _CUDA, + KernelBackend.CUDA_JIT: _CUDA, + } + format_signature = FormatSignature( + supported_dtypes=_NORM_DTYPES, + description="out = (x / RMS(x)) * weight; returns tensor", + ) + descriptions = { + KernelBackend.CUDA_AOT: "RMS normalization (sgl_kernel wheel).", + KernelBackend.CUDA_JIT: "RMS normalization (sglang.jit_kernel).", + KernelBackend.TORCH: "RMS normalization (pure-torch reference).", + } + + def forward_native( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import torch + + x = input.to(torch.float32) + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + result = (x * weight).to(input.dtype) + if out is None: + return result + out.copy_(result) + return out + + def forward_cuda_aot( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import sgl_kernel + + return sgl_kernel.rmsnorm(input, weight, eps, out, enable_pdl) + + def forward_cuda_jit( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import torch + + from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm + + if out is None: + out = torch.empty_like(input) + jit_rmsnorm(input, weight, out, eps) + return out + + +class FusedAddRMSNormOp(BaseFusedOp): + """In-place ``residual += input; input = RMSNorm(residual) * weight``. + + Writes the sum into ``residual`` and the normalized value into ``input``; + returns ``None``. ``enable_pdl`` is honored by the AOT backend only. + """ + + op = "layernorm.fused_add_rmsnorm" + priority = _NORM_PRIORITY + capabilities = { + KernelBackend.CUDA_AOT: _CUDA, + KernelBackend.CUDA_JIT: _CUDA, + } + format_signature = FormatSignature( + supported_dtypes=_NORM_DTYPES, + in_place=True, + description="residual += x; x = RMSNorm(residual) * weight", + ) + descriptions = { + KernelBackend.CUDA_AOT: ( + "Fused residual-add + RMS normalization (sgl_kernel wheel)." + ), + KernelBackend.CUDA_JIT: ( + "Fused residual-add + RMS normalization (sglang.jit_kernel)." + ), + KernelBackend.TORCH: ( + "Fused residual-add + RMS normalization (pure-torch reference)." + ), + } + + def forward_native( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + import torch + + acc = input.to(torch.float32) + residual.to(torch.float32) + residual.copy_(acc.to(residual.dtype)) + variance = acc.pow(2).mean(dim=-1, keepdim=True) + normed = acc * torch.rsqrt(variance + eps) + input.copy_((normed * weight).to(input.dtype)) + + def forward_cuda_aot( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + import sgl_kernel + + return sgl_kernel.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl) + + def forward_cuda_jit( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm + + return jit_fused_add_rmsnorm(input, residual, weight, eps) + + +class GemmaRMSNormOp(BaseFusedOp): + """``out = (input / RMS(input)) * (weight + 1)``; returns a tensor.""" + + op = "layernorm.gemma_rmsnorm" + priority = _NORM_PRIORITY + capabilities = {KernelBackend.CUDA_AOT: _CUDA} + format_signature = FormatSignature( + supported_dtypes=_NORM_DTYPES, + description="out = (x / RMS(x)) * (weight + 1); returns tensor", + ) + descriptions = { + KernelBackend.CUDA_AOT: "Gemma-style RMS normalization (sgl_kernel wheel).", + KernelBackend.TORCH: "Gemma-style RMS normalization (pure-torch reference).", + } + + def forward_native( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import torch + + x = input.to(torch.float32) + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + result = (x * (1.0 + weight.to(torch.float32))).to(input.dtype) + if out is None: + return result + out.copy_(result) + return out + + def forward_cuda_aot( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import sgl_kernel + + return sgl_kernel.gemma_rmsnorm(input, weight, eps, out, enable_pdl) + + +class GemmaFusedAddRMSNormOp(BaseFusedOp): + """In-place ``residual += input; input = GemmaRMSNorm(residual) * (weight + 1)``.""" + + op = "layernorm.gemma_fused_add_rmsnorm" + priority = _NORM_PRIORITY + capabilities = {KernelBackend.CUDA_AOT: _CUDA} + format_signature = FormatSignature( + supported_dtypes=_NORM_DTYPES, + in_place=True, + description="residual += x; x = GemmaRMSNorm(residual) * (weight + 1)", + ) + descriptions = { + KernelBackend.CUDA_AOT: ("Gemma-style fused residual-add + RMS normalization."), + KernelBackend.TORCH: ( + "Gemma-style fused residual-add + RMS normalization " + "(pure-torch reference)." + ), + } + + def forward_native( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + import torch + + acc = input.to(torch.float32) + residual.to(torch.float32) + residual.copy_(acc.to(residual.dtype)) + variance = acc.pow(2).mean(dim=-1, keepdim=True) + normed = acc * torch.rsqrt(variance + eps) + input.copy_((normed * (1.0 + weight.to(torch.float32))).to(input.dtype)) + + def forward_cuda_aot( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + import sgl_kernel + + return sgl_kernel.gemma_fused_add_rmsnorm( + input, residual, weight, eps, enable_pdl + ) + + +_RMSNORM = register_fused_op(RMSNormOp(), __name__, "_RMSNORM") +_FUSED_ADD_RMSNORM = register_fused_op( + FusedAddRMSNormOp(), __name__, "_FUSED_ADD_RMSNORM" +) +_GEMMA_RMSNORM = register_fused_op(GemmaRMSNormOp(), __name__, "_GEMMA_RMSNORM") +_GEMMA_FUSED_ADD_RMSNORM = register_fused_op( + GemmaFusedAddRMSNormOp(), __name__, "_GEMMA_FUSED_ADD_RMSNORM" +) + + +def rmsnorm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, +) -> torch.Tensor: + """RMS normalization: ``out = (input / RMS(input)) * weight``.""" + return _RMSNORM(input, weight, eps, out, enable_pdl) + + +def fused_add_rmsnorm( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, +) -> None: + """In-place fused residual add + RMS normalization.""" + return _FUSED_ADD_RMSNORM(input, residual, weight, eps, enable_pdl) + + +def gemma_rmsnorm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, +) -> torch.Tensor: + """Gemma-style RMS normalization: ``out = (input / RMS(input)) * (weight + 1)``.""" + return _GEMMA_RMSNORM(input, weight, eps, out, enable_pdl) + + +def gemma_fused_add_rmsnorm( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, +) -> None: + """In-place Gemma-style fused residual add + RMS normalization.""" + return _GEMMA_FUSED_ADD_RMSNORM(input, residual, weight, eps, enable_pdl) + + +__all__ = [ + "RMSNormOp", + "FusedAddRMSNormOp", + "GemmaRMSNormOp", + "GemmaFusedAddRMSNormOp", + "rmsnorm", + "fused_add_rmsnorm", + "gemma_rmsnorm", + "gemma_fused_add_rmsnorm", +] diff --git a/python/sglang/kernels/ops/mamba/__init__.py b/python/sglang/kernels/ops/mamba/__init__.py new file mode 100644 index 000000000..eff57abe1 --- /dev/null +++ b/python/sglang/kernels/ops/mamba/__init__.py @@ -0,0 +1,86 @@ +"""State-space / Mamba kernels (causal conv1d).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec + +if TYPE_CHECKING: + import torch + +register_kernel( + KernelSpec( + op="mamba.causal_conv1d_fwd", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel.mamba:causal_conv1d_fwd", + format_signature=FormatSignature( + in_place=True, description="causal depthwise conv1d forward (prefill)" + ), + description="Causal conv1d forward (sgl_kernel wheel).", + ) +) +register_kernel( + KernelSpec( + op="mamba.causal_conv1d_update", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel.mamba:causal_conv1d_update", + format_signature=FormatSignature( + in_place=True, description="causal depthwise conv1d update (decode)" + ), + description="Causal conv1d update (sgl_kernel wheel).", + ) +) + + +def causal_conv1d_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias_: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + silu_activation: bool, + pad_slot_id: int, +): + """Causal depthwise conv1d forward (prefill).""" + return get_kernel("mamba.causal_conv1d_fwd", KernelBackend.CUDA_AOT)( + x, + weight, + bias_, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + silu_activation, + pad_slot_id, + ) + + +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias_: Optional[torch.Tensor], + silu_activation: bool, + cache_seqlens: Optional[torch.Tensor], + conv_state_indices: Optional[torch.Tensor], + pad_slot_id: int, +): + """Causal depthwise conv1d update (decode).""" + return get_kernel("mamba.causal_conv1d_update", KernelBackend.CUDA_AOT)( + x, + conv_state, + weight, + bias_, + silu_activation, + cache_seqlens, + conv_state_indices, + pad_slot_id, + ) + + +__all__ = ["causal_conv1d_fwd", "causal_conv1d_update"] diff --git a/python/sglang/kernels/ops/memory/__init__.py b/python/sglang/kernels/ops/memory/__init__.py new file mode 100644 index 000000000..37f2967c6 --- /dev/null +++ b/python/sglang/kernels/ops/memory/__init__.py @@ -0,0 +1,30 @@ +"""Memory / KV-slot allocation kernels (Triton). + +The Triton kernels migrated here live in this package +(``sglang.kernels.ops.memory.``); import them from there. Their +``KernelSpec`` metadata is registered below for inventory (backend = Triton). +""" + +from sglang.kernels.registry import register_kernel +from sglang.kernels.spec import KernelBackend, KernelSpec + +# (module, public_fn) migrated from mem_cache/triton_ops. +_TRITON_KERNELS = [ + ("allocator", "alloc_extend_kernel"), + ("allocator", "alloc_decode_kernel"), + ("common", "write_req_to_token_pool_triton"), + ("common", "get_last_loc_triton"), + ("common", "get_last_loc_triton_safe"), + ("virtual_slot", "alloc_bind_inplace"), +] +for _mod, _fn in _TRITON_KERNELS: + register_kernel( + KernelSpec( + op=f"memory.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.memory.{_mod}:{_fn}", + ) + ) +del _mod, _fn + +__all__ = [] diff --git a/python/sglang/srt/mem_cache/triton_ops/allocator.py b/python/sglang/kernels/ops/memory/allocator.py similarity index 100% rename from python/sglang/srt/mem_cache/triton_ops/allocator.py rename to python/sglang/kernels/ops/memory/allocator.py diff --git a/python/sglang/srt/mem_cache/triton_ops/common.py b/python/sglang/kernels/ops/memory/common.py similarity index 100% rename from python/sglang/srt/mem_cache/triton_ops/common.py rename to python/sglang/kernels/ops/memory/common.py diff --git a/python/sglang/srt/mem_cache/triton_ops/virtual_slot.py b/python/sglang/kernels/ops/memory/virtual_slot.py similarity index 100% rename from python/sglang/srt/mem_cache/triton_ops/virtual_slot.py rename to python/sglang/kernels/ops/memory/virtual_slot.py diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py new file mode 100644 index 000000000..45e70bcbe --- /dev/null +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -0,0 +1,119 @@ +"""Mixture-of-Experts routing / bookkeeping kernels.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, +) + +if TYPE_CHECKING: + import torch + +_CUDA = CapabilityRequirement(requires_cuda=True) + +register_kernel( + KernelSpec( + op="moe.moe_align_block_size", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel:moe_align_block_size", + format_signature=FormatSignature( + in_place=True, + description="align/sort expert token ids into block-padded buffers", + ), + description="MoE align-block-size (sgl_kernel wheel).", + ) +) +register_kernel( + KernelSpec( + op="moe.moe_align_block_size", + backend=KernelBackend.CUDA_JIT, + target="sglang.jit_kernel.moe_align:moe_align_block_size", + capability=_CUDA, + format_signature=FormatSignature( + in_place=True, + description="MoE align-block-size (JIT variant, AOT signature)", + ), + description="MoE align-block-size (sglang.jit_kernel).", + ) +) +register_kernel( + KernelSpec( + op="moe.topk_softmax", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel:topk_softmax", + format_signature=FormatSignature( + in_place=True, + description="top-k softmax routing weights/ids", + ), + description="MoE top-k softmax (sgl_kernel wheel).", + ) +) + + +def moe_align_block_size( + topk_ids: torch.Tensor, + num_experts: int, + block_size: int, + sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + cumsum_buffer: torch.Tensor, + pad_sorted_token_ids: bool = False, +) -> None: + """Align and sort expert token ids into block-padded output buffers.""" + return get_kernel("moe.moe_align_block_size", KernelBackend.CUDA_AOT)( + topk_ids, + num_experts, + block_size, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + cumsum_buffer, + pad_sorted_token_ids, + ) + + +def topk_softmax( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + moe_softcapping: float = 0.0, + correction_bias: Optional[torch.Tensor] = None, +) -> None: + """Compute top-k softmax routing weights/ids for MoE.""" + return get_kernel("moe.topk_softmax", KernelBackend.CUDA_AOT)( + topk_weights, + topk_ids, + gating_output, + renormalize, + moe_softcapping, + correction_bias, + ) + + +__all__ = ["moe_align_block_size", "topk_softmax"] + + +# Fused MoE-LoRA Triton kernels migrated into this group (from lora/triton_ops); +# registered for inventory. Import them from their modules. +_TRITON_KERNELS = [ + ("fused_moe_lora_kernel", "fused_moe_lora"), + ("virtual_experts", "merged_experts_fused_moe_lora_add"), +] +for _mod, _fn in _TRITON_KERNELS: + register_kernel( + KernelSpec( + op=f"moe.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.moe.{_mod}:{_fn}", + ) + ) +del _mod, _fn diff --git a/python/sglang/srt/lora/triton_ops/fused_moe_lora_kernel.py b/python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py similarity index 100% rename from python/sglang/srt/lora/triton_ops/fused_moe_lora_kernel.py rename to python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py diff --git a/python/sglang/kernels/ops/moe/trtllm_lora_temp/__init__.py b/python/sglang/kernels/ops/moe/trtllm_lora_temp/__init__.py new file mode 100644 index 000000000..fdc403e21 --- /dev/null +++ b/python/sglang/kernels/ops/moe/trtllm_lora_temp/__init__.py @@ -0,0 +1,3 @@ +"""Experimental TRT-LLM LoRA kernel variants (gated by ``SGLANG_EXPERIMENTAL_LORA_OPTI`` / ``lora_envs``). + +Migrated from ``sglang.srt.lora.trtllm_lora_temp.triton_ops`` (RFC #29630).""" diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py b/python/sglang/kernels/ops/moe/trtllm_lora_temp/virtual_experts.py similarity index 99% rename from python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py rename to python/sglang/kernels/ops/moe/trtllm_lora_temp/virtual_experts.py index bf184f2be..6b26f3dd4 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py +++ b/python/sglang/kernels/ops/moe/trtllm_lora_temp/virtual_experts.py @@ -10,10 +10,10 @@ import triton import triton.language as tl from sglang.jit_kernel.moe_align import moe_align_block_size as jit_moe_align_block_size -from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs -from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import ( +from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import ( get_pdl_launch_metadata, ) +from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs @triton.jit diff --git a/python/sglang/srt/lora/triton_ops/virtual_experts.py b/python/sglang/kernels/ops/moe/virtual_experts.py similarity index 100% rename from python/sglang/srt/lora/triton_ops/virtual_experts.py rename to python/sglang/kernels/ops/moe/virtual_experts.py diff --git a/python/sglang/kernels/ops/quantization/__init__.py b/python/sglang/kernels/ops/quantization/__init__.py new file mode 100644 index 000000000..0159dd386 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/__init__.py @@ -0,0 +1,122 @@ +"""Quantization kernels (per-token / per-token-group FP8 & INT8).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, +) + +if TYPE_CHECKING: + import torch + +_CUDA = CapabilityRequirement(requires_cuda=True) + +register_kernel( + KernelSpec( + op="quantization.sgl_per_token_quant_fp8", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel:sgl_per_token_quant_fp8", + format_signature=FormatSignature( + supported_dtypes=("float8_e4m3fn",), + in_place=True, + description="per-token FP8 quantization into output_q/output_s", + ), + description="Per-token FP8 quantization (sgl_kernel wheel).", + ) +) +# fp8 / int8 are legacy aliases of the same 8bit kernel in the wheel; register +# each public name so runtime imports resolve to a stable spec. +for _name in ( + "sgl_per_token_group_quant_8bit", + "sgl_per_token_group_quant_fp8", + "sgl_per_token_group_quant_int8", +): + register_kernel( + KernelSpec( + op=f"quantization.{_name}", + backend=KernelBackend.CUDA_AOT, + target=f"sgl_kernel:{_name}", + format_signature=FormatSignature( + in_place=True, + description="per-token-group 8-bit quantization", + ), + description=f"{_name} (sgl_kernel wheel).", + ) + ) +del _name + +register_kernel( + KernelSpec( + op="quantization.sgl_per_token_group_quant_8bit", + backend=KernelBackend.CUDA_JIT, + target="sglang.jit_kernel.per_token_group_quant_8bit:per_token_group_quant_8bit", + capability=_CUDA, + format_signature=FormatSignature( + in_place=True, + description="per-token-group 8-bit quantization (JIT variant)", + ), + description="Per-token-group 8-bit quantization (sglang.jit_kernel).", + ) +) + + +def sgl_per_token_quant_fp8( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, +) -> None: + """Per-token FP8 quantization, writing into ``output_q`` / ``output_s``.""" + return get_kernel("quantization.sgl_per_token_quant_fp8", KernelBackend.CUDA_AOT)( + input, output_q, output_s + ) + + +def sgl_per_token_group_quant_8bit( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + scale_ue8m0: bool = False, + fuse_silu_and_mul: bool = False, + masked_m: Optional[torch.Tensor] = None, + enable_v2: Optional[bool] = None, +) -> None: + """Per-token-group 8-bit quantization, writing into ``output_q`` / ``output_s``.""" + return get_kernel( + "quantization.sgl_per_token_group_quant_8bit", KernelBackend.CUDA_AOT + )( + input, + output_q, + output_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + fuse_silu_and_mul, + masked_m, + enable_v2, + ) + + +# Legacy aliases kept for source compatibility with existing call sites. +sgl_per_token_group_quant_fp8 = sgl_per_token_group_quant_8bit +sgl_per_token_group_quant_int8 = sgl_per_token_group_quant_8bit + + +__all__ = [ + "sgl_per_token_quant_fp8", + "sgl_per_token_group_quant_8bit", + "sgl_per_token_group_quant_fp8", + "sgl_per_token_group_quant_int8", +] diff --git a/python/sglang/kernels/ops/sampling/__init__.py b/python/sglang/kernels/ops/sampling/__init__.py new file mode 100644 index 000000000..8f9587e58 --- /dev/null +++ b/python/sglang/kernels/ops/sampling/__init__.py @@ -0,0 +1,56 @@ +"""Sampling kernels (top-k / top-p probability renormalization).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Union + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec + +if TYPE_CHECKING: + import torch + +register_kernel( + KernelSpec( + op="sampling.top_k_renorm_probs", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel.sampling:top_k_renorm_probs", + format_signature=FormatSignature( + description="renormalize probs by top-k thresholding; returns tensor" + ), + description="Top-k probability renormalization (sgl_kernel wheel).", + ) +) +register_kernel( + KernelSpec( + op="sampling.top_p_renorm_probs", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel.sampling:top_p_renorm_probs", + format_signature=FormatSignature( + description="renormalize probs by top-p thresholding; returns tensor" + ), + description="Top-p probability renormalization (sgl_kernel wheel).", + ) +) + + +def top_k_renorm_probs( + probs: torch.Tensor, top_k: Union[torch.Tensor, int] +) -> torch.Tensor: + """Renormalize ``probs`` by top-k thresholding.""" + return get_kernel("sampling.top_k_renorm_probs", KernelBackend.CUDA_AOT)( + probs, top_k + ) + + +def top_p_renorm_probs( + probs: torch.Tensor, top_p: Union[torch.Tensor, float] +) -> torch.Tensor: + """Renormalize ``probs`` by top-p thresholding.""" + return get_kernel("sampling.top_p_renorm_probs", KernelBackend.CUDA_AOT)( + probs, top_p + ) + + +__all__ = ["top_k_renorm_probs", "top_p_renorm_probs"] diff --git a/python/sglang/kernels/ops/spatial/__init__.py b/python/sglang/kernels/ops/spatial/__init__.py new file mode 100644 index 000000000..9542f244a --- /dev/null +++ b/python/sglang/kernels/ops/spatial/__init__.py @@ -0,0 +1,49 @@ +"""Spatial / green-context stream helpers.""" + +from __future__ import annotations + +from typing import Optional + +from sglang.kernels.registry import register_kernel +from sglang.kernels.selector import get_kernel +from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec + +register_kernel( + KernelSpec( + op="spatial.get_sm_available", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel.spatial:get_sm_available", + format_signature=FormatSignature( + description="number of SMs available on device" + ), + description="Query available SM count (sgl_kernel wheel).", + ) +) +register_kernel( + KernelSpec( + op="spatial.create_greenctx_stream_by_value", + backend=KernelBackend.CUDA_AOT, + target="sgl_kernel.spatial:create_greenctx_stream_by_value", + format_signature=FormatSignature( + description="create two green-context streams partitioned by SM count" + ), + description="Green-context stream creation (sgl_kernel wheel).", + ) +) + + +def get_sm_available(device_id: Optional[int] = None) -> int: + """Return the number of SMs available on ``device_id``.""" + return get_kernel("spatial.get_sm_available", KernelBackend.CUDA_AOT)(device_id) + + +def create_greenctx_stream_by_value( + SM_a: int, SM_b: int, device_id: Optional[int] = None +): + """Create two green-context streams partitioned by ``SM_a`` / ``SM_b``.""" + return get_kernel( + "spatial.create_greenctx_stream_by_value", KernelBackend.CUDA_AOT + )(SM_a, SM_b, device_id) + + +__all__ = ["get_sm_available", "create_greenctx_stream_by_value"] diff --git a/python/sglang/kernels/ops/speculative/__init__.py b/python/sglang/kernels/ops/speculative/__init__.py new file mode 100644 index 000000000..209f12cc7 --- /dev/null +++ b/python/sglang/kernels/ops/speculative/__init__.py @@ -0,0 +1,33 @@ +"""Speculative-decoding kernels (Triton). + +The Triton kernels migrated here live in this package +(``sglang.kernels.ops.speculative.``); import them from there. Their +``KernelSpec`` metadata is registered below for inventory (backend = Triton). +""" + +from sglang.kernels.registry import register_kernel +from sglang.kernels.spec import KernelBackend, KernelSpec + +# (module, public_fn) migrated from speculative/triton_ops. +_TRITON_KERNELS = [ + ("cache_locs", "assign_req_to_token_pool_func"), + ("cache_locs", "assign_extend_cache_locs_func"), + ("cache_locs", "generate_draft_decode_kv_indices"), + ("eagle", "fill_bonus_tokens"), + ("eagle", "fill_accept_out_cache_loc"), + ("gather_spec_extras", "gather_spec_extras"), + ("multi_layer_eagle", "rotate_input_ids_triton"), + ("spec_tree", "sgl_build_tree_kernel_efficient_triton"), + ("spec_tree", "verify_tree_greedy_kernel_triton"), +] +for _mod, _fn in _TRITON_KERNELS: + register_kernel( + KernelSpec( + op=f"speculative.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.speculative.{_mod}:{_fn}", + ) + ) +del _mod, _fn + +__all__ = [] diff --git a/python/sglang/srt/speculative/triton_ops/cache_locs.py b/python/sglang/kernels/ops/speculative/cache_locs.py similarity index 100% rename from python/sglang/srt/speculative/triton_ops/cache_locs.py rename to python/sglang/kernels/ops/speculative/cache_locs.py diff --git a/python/sglang/srt/speculative/triton_ops/dflash.py b/python/sglang/kernels/ops/speculative/dflash.py similarity index 100% rename from python/sglang/srt/speculative/triton_ops/dflash.py rename to python/sglang/kernels/ops/speculative/dflash.py diff --git a/python/sglang/srt/speculative/triton_ops/eagle.py b/python/sglang/kernels/ops/speculative/eagle.py similarity index 100% rename from python/sglang/srt/speculative/triton_ops/eagle.py rename to python/sglang/kernels/ops/speculative/eagle.py diff --git a/python/sglang/srt/speculative/triton_ops/fused_kv_materialize.py b/python/sglang/kernels/ops/speculative/fused_kv_materialize.py similarity index 100% rename from python/sglang/srt/speculative/triton_ops/fused_kv_materialize.py rename to python/sglang/kernels/ops/speculative/fused_kv_materialize.py diff --git a/python/sglang/srt/speculative/triton_ops/gather_spec_extras.py b/python/sglang/kernels/ops/speculative/gather_spec_extras.py similarity index 100% rename from python/sglang/srt/speculative/triton_ops/gather_spec_extras.py rename to python/sglang/kernels/ops/speculative/gather_spec_extras.py diff --git a/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py b/python/sglang/kernels/ops/speculative/multi_layer_eagle.py similarity index 100% rename from python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py rename to python/sglang/kernels/ops/speculative/multi_layer_eagle.py diff --git a/python/sglang/srt/speculative/triton_ops/spec_tree.py b/python/sglang/kernels/ops/speculative/spec_tree.py similarity index 100% rename from python/sglang/srt/speculative/triton_ops/spec_tree.py rename to python/sglang/kernels/ops/speculative/spec_tree.py diff --git a/python/sglang/kernels/registry.py b/python/sglang/kernels/registry.py new file mode 100644 index 000000000..bbf8b3bf1 --- /dev/null +++ b/python/sglang/kernels/registry.py @@ -0,0 +1,71 @@ +"""In-memory registry of :class:`KernelSpec` entries. + +The registry is the single inventory of "which operators have which backend +implementations". It is populated at import time by the ``sglang.kernels.ops.*`` +group packages, using only metadata (import path strings) — registering a spec +never imports ``torch`` or a kernel backend. +""" + +from __future__ import annotations + +from collections import defaultdict +from typing import Dict, List + +from sglang.kernels.spec import KernelBackend, KernelSpec + + +class KernelRegistry: + """Maps ``"."`` operator ids to their :class:`KernelSpec` list.""" + + def __init__(self) -> None: + self._by_op: Dict[str, List[KernelSpec]] = defaultdict(list) + + def register(self, spec: KernelSpec) -> KernelSpec: + """Register ``spec``. + + Re-registering the same ``(op, backend)`` pair replaces the previous + entry so that module reloads during tests stay idempotent. + """ + existing = self._by_op[spec.op] + for i, other in enumerate(existing): + if other.backend == spec.backend: + existing[i] = spec + return spec + existing.append(spec) + return spec + + def get(self, op: str) -> List[KernelSpec]: + """All registered specs for ``op`` (empty list if none).""" + return list(self._by_op.get(op, ())) + + def get_backend(self, op: str, backend: KernelBackend) -> KernelSpec: + """The spec for ``op`` provided by ``backend``. + + Raises ``KeyError`` if no such implementation is registered. + """ + for spec in self._by_op.get(op, ()): + if spec.backend == backend: + return spec + raise KeyError(f"No '{backend.value}' backend registered for op {op!r}") + + def has(self, op: str) -> bool: + return bool(self._by_op.get(op)) + + def ops(self) -> List[str]: + """Sorted list of all registered operator ids.""" + return sorted(self._by_op.keys()) + + def all_specs(self) -> List[KernelSpec]: + specs: List[KernelSpec] = [] + for op in self.ops(): + specs.extend(self._by_op[op]) + return specs + + +# Process-wide registry. Group packages register into this instance on import. +registry = KernelRegistry() + + +def register_kernel(spec: KernelSpec) -> KernelSpec: + """Register ``spec`` in the process-wide :data:`registry`.""" + return registry.register(spec) diff --git a/python/sglang/kernels/selector.py b/python/sglang/kernels/selector.py new file mode 100644 index 000000000..472d84232 --- /dev/null +++ b/python/sglang/kernels/selector.py @@ -0,0 +1,78 @@ +"""Fixed-path kernel resolution over the :data:`registry`. + +There is no priority ranking or heuristic backend selection. Each operator has +a fixed call path — its :attr:`KernelSpec.target`: + +- an op with a single registered backend resolves to it directly; +- an op with several registered backends must be resolved by naming the backend + explicitly (``backend=...``). The extra backends exist only as inventory, and + are never silently auto-picked. + +:func:`get_kernel` is the fast path used by the public ``ops.*`` wrappers: it +resolves the spec to its callable and caches the result so repeated calls do +not re-run resolution or re-import. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Callable, Optional + +from sglang.kernels.registry import registry +from sglang.kernels.spec import KernelBackend, KernelSpec + + +def select_kernel(op: str, backend: Optional[KernelBackend] = None) -> KernelSpec: + """Return the :class:`KernelSpec` for ``op`` (its fixed call path). + + Parameters + ---------- + op: + Operator id, ``"."``. + backend: + Required only when ``op`` has more than one registered backend; selects + which one. For single-backend ops it is optional. + + Raises + ------ + KeyError + If ``op`` is unknown, or if ``backend`` is requested but not registered. + ValueError + If ``op`` has multiple backends and ``backend`` is not given. + """ + specs = registry.get(op) + if not specs: + raise KeyError(f"No kernels registered for op {op!r}") + + if backend is not None: + for spec in specs: + if spec.backend == backend: + return spec + raise KeyError(f"No '{backend.value}' backend registered for op {op!r}") + + if len(specs) == 1: + return specs[0] + + raise ValueError( + f"op {op!r} has multiple registered backends " + f"({[s.backend.value for s in specs]}); pass backend=... to choose one" + ) + + +@lru_cache(maxsize=None) +def _resolve(op: str, backend: Optional[KernelBackend]) -> Callable: + return select_kernel(op, backend=backend).load() + + +def get_kernel(op: str, backend: Optional[KernelBackend] = None) -> Callable: + """Resolve ``op`` to a callable kernel and cache it. + + This is what the public ``sglang.kernels.ops.*`` wrappers call. The first + call resolves and imports the backend; later calls hit the cache. + """ + return _resolve(op, backend) + + +def clear_cache() -> None: + """Drop the resolved-callable cache (used by tests).""" + _resolve.cache_clear() diff --git a/python/sglang/kernels/spec.py b/python/sglang/kernels/spec.py new file mode 100644 index 000000000..c854824a2 --- /dev/null +++ b/python/sglang/kernels/spec.py @@ -0,0 +1,185 @@ +"""Lightweight metadata for the unified ``sglang.kernels`` namespace. + +This module defines small, dependency-free descriptors used to *inventory* +kernel implementations and drive a simple, heuristic dispatch. It intentionally +does not import ``torch``, ``sgl_kernel`` or ``sglang.jit_kernel`` at module +import time so that ``import sglang.kernels`` stays cheap and works on a CPU-only +box (see RFC #29630, Phase 2). + +The concrete callable behind a :class:`KernelSpec` is resolved lazily through +``KernelSpec.load()``; nothing is imported until a kernel is actually called. +""" + +from __future__ import annotations + +import importlib +from enum import Enum +from typing import Callable, Optional, Tuple + +import msgspec + + +class KernelBackend(str, Enum): + """Implementation backend for a kernel. + + Values mirror the backends called out in RFC #29630: JIT CUDA, AOT + CUDA/C++ (the ``sgl_kernel`` wheel), Triton, CuTe DSL, FlashInfer, DeepGEMM, + and the pure-``torch`` fallback path. + """ + + TORCH = "torch" # pure-torch reference (forward_native) + TORCH_COMPILE = "torch_compile" # torch.compile(forward_native) + TRITON = "triton" + CUDA_JIT = "cuda_jit" # sglang.jit_kernel + CUDA_AOT = "cuda_aot" # sgl_kernel wheel + CUTE_DSL = "cute_dsl" + FLASHINFER = "flashinfer" + DEEPGEMM = "deepgemm" + # TODO(RFC #29630): backends for other hardware (hip_c / npu / cpu-avx, ...) + + +class PlatformInfo(msgspec.Struct, frozen=True): + """A minimal snapshot of the runtime accelerator platform. + + Kept torch-free at import time; use :meth:`detect` to build one from the + live process (which does import ``torch``). + """ + + device_type: str = "cpu" # "cuda", "hip", "cpu", ... + cuda_arch_major: Optional[int] = None + cuda_arch_minor: Optional[int] = None + + @property + def is_cuda(self) -> bool: + return self.device_type == "cuda" + + @property + def is_hip(self) -> bool: + return self.device_type == "hip" + + @classmethod + def detect(cls) -> PlatformInfo: + """Build a :class:`PlatformInfo` from the current process. + + Never raises: if ``torch`` is missing or no accelerator is visible the + default CPU platform is returned. + """ + try: + import torch + except Exception: + return cls() + + try: + if torch.version.hip is not None and torch.cuda.is_available(): + return cls(device_type="hip") + if torch.cuda.is_available(): + major, minor = torch.cuda.get_device_capability() + return cls( + device_type="cuda", + cuda_arch_major=major, + cuda_arch_minor=minor, + ) + except Exception: + pass + return cls() + + +class CapabilityRequirement(msgspec.Struct, frozen=True): + """Coarse hardware requirement used to filter out unusable backends. + + ``min_cuda_arch`` / ``max_cuda_arch`` are ``(major, minor)`` tuples, e.g. + ``(9, 0)`` for SM90. They only apply when the kernel requires CUDA. + """ + + requires_cuda: bool = False + requires_hip: bool = False + min_cuda_arch: Optional[Tuple[int, int]] = None + max_cuda_arch: Optional[Tuple[int, int]] = None + + def is_satisfied_by(self, platform: PlatformInfo) -> bool: + if self.requires_hip and not platform.is_hip: + return False + if self.requires_cuda and not platform.is_cuda: + return False + if platform.is_cuda and platform.cuda_arch_major is not None: + arch = (platform.cuda_arch_major, platform.cuda_arch_minor or 0) + if self.min_cuda_arch is not None and arch < self.min_cuda_arch: + return False + if self.max_cuda_arch is not None and arch > self.max_cuda_arch: + return False + return True + + +class FormatSignature(msgspec.Struct, frozen=True): + """A light description of a kernel's data contract. + + This is deliberately loose in the first version — enough to document intent + and support future inventory tooling, not a strict schema. + """ + + supported_dtypes: Tuple[str, ...] = () + in_place: bool = False + description: str = "" + + +class KernelSpec(msgspec.Struct, frozen=True): + """A single callable kernel implementation and its metadata. + + Parameters + ---------- + op: + Fully-qualified operator id, ``"."`` (e.g. + ``"layernorm.rmsnorm"``). This is the public lookup key. + backend: + Which :class:`KernelBackend` provides this implementation. + target: + Import path of the callable in ``"module:attr"`` form, resolved lazily + by :meth:`load` (e.g. ``"sgl_kernel:rmsnorm"``). ``attr`` may be a + dotted path into a module-level object, e.g. + ``"sglang.kernels.ops.layernorm:_RMSNORM.forward_cuda_aot"`` for a + bound :class:`~sglang.kernels.fused_op.BaseFusedOp` backend method. + capability: + Hardware requirement used by the selector to skip unusable backends. + format_signature: + Optional data-contract description for inventory/documentation. + description: + Human-readable one-liner. + """ + + op: str + backend: KernelBackend + target: str + capability: CapabilityRequirement = msgspec.field( + default_factory=CapabilityRequirement + ) + format_signature: FormatSignature = msgspec.field(default_factory=FormatSignature) + description: str = "" + + @property + def group(self) -> str: + return self.op.split(".", 1)[0] + + @property + def name(self) -> str: + return self.op.split(".", 1)[1] if "." in self.op else self.op + + def is_available(self, platform: PlatformInfo) -> bool: + """Whether this backend can run on ``platform`` (metadata-only check).""" + return self.capability.is_satisfied_by(platform) + + def load(self) -> Callable: + """Import and return the backing callable. + + Raises the underlying ``ImportError`` / ``AttributeError`` if the + backend is not installed on this platform — call sites decide how to + handle that. + """ + module_path, sep, attr = self.target.partition(":") + if not sep or not attr: + raise ValueError( + f"KernelSpec.target must be 'module:attr', got {self.target!r}" + ) + obj = importlib.import_module(module_path) + for part in attr.split("."): + obj = getattr(obj, part) + return obj diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index c02021188..19fdc4bbf 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -43,14 +43,14 @@ _is_hip = is_hip() if _is_hip: from sgl_kernel import apply_token_bitmask_inplace_cuda else: - from sglang.srt.constrained.triton_ops.bitmask_ops import ( + from sglang.kernels.ops.grammar.bitmask_ops import ( apply_token_bitmask_inplace_triton, ) +from sglang.kernels.ops.grammar.token_filter_ops import set_token_filter_triton from sglang.srt.constrained.torch_ops.token_filter_torch_ops import ( set_token_filter_torch, ) -from sglang.srt.constrained.triton_ops.token_filter_ops import set_token_filter_triton logger = logging.getLogger(__name__) MAX_ROLLBACK_TOKENS = 200 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 3a724e7d6..2b3ad5cc5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -648,6 +648,8 @@ class Envs: SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False) SGLANG_DSA_TOPK_BROADCAST = EnvBool(False) SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False) + SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True) + SGLANG_DSA_USE_FUSED_METADATA_GENERATION = EnvBool(True) # sgl-kernel SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False) @@ -656,6 +658,11 @@ class Envs: SGLANG_USE_SGL_FA3_KERNEL = EnvBool(True) # Kernels + # Force every sglang.kernels BaseFusedOp onto one backend (a KernelBackend + # value, e.g. "torch" / "torch_compile" / "triton" / "cuda_aot"); unset = + # auto-select by priority. "torch" flips all fused ops to their pure-torch + # reference implementations for numerical-bug bisection. + SGLANG_FORCE_FUSED_OP_BACKEND = EnvStr(None) USE_TRITON_W8A8_FP8_KERNEL = EnvBool(False) SGLANG_RETURN_ORIGINAL_LOGPROB = EnvBool(False) SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 12fd629cc..6ab02375a 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -14,11 +14,11 @@ from typing import TYPE_CHECKING, Optional import torch import triton -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.triton_ops.aiter_unified_attention import ( +from sglang.kernels.ops.kvcache.aiter_unified_attention import ( scatter_ragged_to_page_table_kernel, scatter_req_to_token_to_page_table_kernel, ) +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.utils import ( assert_buffer_fits, create_flashinfer_kv_indices_triton, @@ -136,7 +136,7 @@ class AiterAttnBackend(AttentionBackend): ): super().__init__() # Lazy import to avoid the initialization of cuda context - from sglang.srt.layers.attention.triton_ops.extend_attention import ( + from sglang.kernels.ops.attention.extend_attention import ( extend_attention_fwd, ) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py index 4b2be7dee..c48126602 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Optional import torch +from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import compute_dsa_seqlens from sglang.srt.layers.attention.utils import seqlens_expand_triton from sglang.srt.utils import is_cuda, is_hip @@ -20,6 +21,9 @@ if TYPE_CHECKING: _is_cuda = is_cuda() _is_hip = is_hip() +_USE_FUSED_METADATA_GENERATION = ( + envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip +) @dataclass @@ -121,8 +125,8 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: """Precompute metadata for normal decode mode.""" max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1] - if _is_cuda and not _is_hip: - from sglang.srt.layers.attention.triton_ops.dsa_metadata import ( + if _USE_FUSED_METADATA_GENERATION and _is_cuda and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_decode_metadata, ) @@ -240,8 +244,8 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: max_seqlen_k = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1] seqlens_expanded_size = bs * self.speculative_num_draft_tokens - if _is_cuda and not _is_hip: - from sglang.srt.layers.attention.triton_ops.dsa_metadata import ( + if _USE_FUSED_METADATA_GENERATION and _is_cuda and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_target_verify_metadata, ) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index c4055f890..5a6467b74 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -138,6 +138,13 @@ def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tens # Reuse this workspace buffer across all DSA backend instances +# Control whether to use fused metadata copy kernel for cuda graph replay (default: enabled) +# Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable +_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip +_USE_FUSED_METADATA_GENERATION = ( + envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip +) + @dataclass(frozen=True) class DSAFlashMLAMetadata: @@ -1374,8 +1381,8 @@ class DeepseekSparseAttnBackend( # Normal Decode max_len = self._graph_page_table_width(metadata) - if is_cuda() and not _is_hip: - from sglang.srt.layers.attention.triton_ops.dsa_metadata import ( + if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_decode_metadata, ) @@ -1416,8 +1423,8 @@ class DeepseekSparseAttnBackend( elif forward_mode.is_target_verify(): max_seqlen_k = self._graph_page_table_width(metadata) - if is_cuda() and not _is_hip: - from sglang.srt.layers.attention.triton_ops.dsa_metadata import ( + if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_target_verify_metadata, ) @@ -1514,8 +1521,8 @@ class DeepseekSparseAttnBackend( device=self.device, ) - if is_cuda() and not _is_hip: - from sglang.srt.layers.attention.triton_ops.dsa_metadata import ( + if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_draft_extend_metadata, ) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index f9dc630fb..94d307d67 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -8,15 +8,15 @@ import torch import triton import triton.language as tl -from sglang.srt.configs.model_config import AttentionArch -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.triton_ops.metadata import ( +from sglang.kernels.ops.attention.metadata import ( normal_decode_set_metadata, prepare_swa_spec_page_table_triton, ) -from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import ( +from sglang.kernels.ops.kvcache.trtllm_mha_page_table import ( build_trtllm_mha_page_table, ) +from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.utils import assert_buffer_fits from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy from sglang.srt.layers.cp.utils import is_cp_v2_active diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index d06e8cda4..8e537259a 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -86,7 +86,7 @@ if is_flashinfer_available(): ) from flashinfer.cascade import merge_state - from sglang.srt.layers.attention.triton_ops.merge_state import merge_state_triton + from sglang.kernels.ops.attention.merge_state import merge_state_triton # FlashInfer's MergeState CUDA kernel uses blockDim = (head_dim/vec_size, num_heads). # When num_heads is large (e.g. with DP attention where attention_tp_size=1), the diff --git a/python/sglang/srt/layers/attention/merge_state.py b/python/sglang/srt/layers/attention/merge_state.py index 245418a9d..ac73b8440 100644 --- a/python/sglang/srt/layers/attention/merge_state.py +++ b/python/sglang/srt/layers/attention/merge_state.py @@ -3,7 +3,7 @@ from typing import Optional, Tuple import torch from sgl_kernel import merge_state_v2 -from sglang.srt.layers.attention.triton_ops.merge_state import merge_state_triton +from sglang.kernels.ops.attention.merge_state import merge_state_triton from sglang.srt.utils import is_cuda _is_cuda = is_cuda() diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 0be715ac1..d7176d261 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -6,6 +6,10 @@ from typing import TYPE_CHECKING, List, Optional import torch import triton +from sglang.kernels.ops.attention.metadata import get_num_kv_splits_triton +from sglang.kernels.ops.kvcache.kv_indices import ( + create_flashinfer_kv_indices_triton, +) from sglang.srt.configs.model_config import AttentionArch from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -13,10 +17,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.parallel_state import get_dcp_group from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.triton_ops.kv_indices import ( - create_flashinfer_kv_indices_triton, -) -from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton from sglang.srt.layers.dcp import ( cp_lse_ag_out_rs_mha, create_triton_kv_indices_for_dcp_triton, @@ -115,15 +115,15 @@ class TritonAttnBackend(AttentionBackend): kv_indptr_buf: Optional[torch.Tensor] = None, ): # Lazy import to avoid the initialization of cuda context - from sglang.srt.layers.attention.triton_ops.decode_attention import ( + from sglang.kernels.ops.attention.decode_attention import ( decode_attention_fwd, ) - from sglang.srt.layers.attention.triton_ops.extend_attention import ( + from sglang.kernels.ops.attention.extend_attention import ( build_unified_kv_indices, extend_attention_fwd, extend_attention_fwd_unified, ) - from sglang.srt.layers.attention.triton_ops.verify_splitkv import ( + from sglang.kernels.ops.attention.verify_splitkv import ( verify_splitkv_fwd, ) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index e19e59141..d58f3f661 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -11,22 +11,22 @@ from typing import TYPE_CHECKING, Optional import torch -from sglang.srt.environ import envs -from sglang.srt.layers.attention.flashinfer_backend import ( - FlashInferAttnBackend, - FlashInferMultiStepDraftBackend, -) -from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import ( +from sglang.kernels.ops.kvcache.trtllm_fp8_kv_kernel import ( fused_fp8_set_kv_buffer, ) -from sglang.srt.layers.attention.triton_ops.trtllm_mha_graph_metadata import ( +from sglang.kernels.ops.kvcache.trtllm_mha_graph_metadata import ( Q_MODE_NONE, Q_MODE_STRIDED, update_trtllm_mha_graph_metadata, ) -from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import ( +from sglang.kernels.ops.kvcache.trtllm_mha_page_table import ( build_trtllm_mha_page_table, ) +from sglang.srt.environ import envs +from sglang.srt.layers.attention.flashinfer_backend import ( + FlashInferAttnBackend, + FlashInferMultiStepDraftBackend, +) from sglang.srt.layers.attention.utils import canonicalize_stride from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 28177c251..125d811ca 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -13,21 +13,21 @@ import torch import triton from sglang.jit_kernel.fixup_zero_kv import fixup_zero_kv_rows -from sglang.srt.environ import envs -from sglang.srt.layers.attention.flashinfer_mla_backend import ( - FlashInferMLAAttnBackend, - FlashInferMLAMultiStepDraftBackend, +from sglang.kernels.ops.attention.pad import ( + pad_draft_extend_query as pad_draft_extend_query_triton, ) -from sglang.srt.layers.attention.triton_ops.kv_indices import ( +from sglang.kernels.ops.attention.pad import ( + unpad_draft_extend_output as unpad_draft_extend_output_triton, +) +from sglang.kernels.ops.kvcache.kv_indices import ( create_flashmla_kv_indices_triton, get_num_kv_index_blocks_flashmla, get_num_page_per_block_flashmla, ) -from sglang.srt.layers.attention.triton_ops.pad import ( - pad_draft_extend_query as pad_draft_extend_query_triton, -) -from sglang.srt.layers.attention.triton_ops.pad import ( - unpad_draft_extend_output as unpad_draft_extend_output_triton, +from sglang.srt.environ import envs +from sglang.srt.layers.attention.flashinfer_mla_backend import ( + FlashInferMLAAttnBackend, + FlashInferMLAMultiStepDraftBackend, ) from sglang.srt.layers.attention.utils import ( concat_mla_absorb_q_general, diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py index d31004baa..6f479c953 100644 --- a/python/sglang/srt/layers/attention/utils.py +++ b/python/sglang/srt/layers/attention/utils.py @@ -3,43 +3,43 @@ import triton import triton.language as tl from sglang.jit_kernel.utils import is_arch_support_pdl -from sglang.srt.layers.attention.triton_ops.cache_ops import ( - concat_and_cast_mha_k_kernel as concat_and_cast_mha_k_kernel, -) -from sglang.srt.layers.attention.triton_ops.cache_ops import ( - concat_and_cast_mha_k_triton as concat_and_cast_mha_k_triton, -) -from sglang.srt.layers.attention.triton_ops.cache_ops import ( - launch_reshape_and_cache_flash as launch_reshape_and_cache_flash, -) -from sglang.srt.layers.attention.triton_ops.cache_ops import ( - reshape_and_cache_flash as reshape_and_cache_flash, -) -from sglang.srt.layers.attention.triton_ops.kv_indices import ( - create_flashinfer_kv_indices_triton as create_flashinfer_kv_indices_triton, -) -from sglang.srt.layers.attention.triton_ops.kv_indices import ( - create_flashmla_kv_indices_triton as create_flashmla_kv_indices_triton, -) -from sglang.srt.layers.attention.triton_ops.kv_indices import ( - get_num_kv_index_blocks_flashmla as get_num_kv_index_blocks_flashmla, -) -from sglang.srt.layers.attention.triton_ops.kv_indices import ( - get_num_page_per_block_flashmla as get_num_page_per_block_flashmla, -) -from sglang.srt.layers.attention.triton_ops.pad import ( +from sglang.kernels.ops.attention.pad import ( pad_sequence_with_mask as pad_sequence_with_mask, ) -from sglang.srt.layers.attention.triton_ops.pad import ( +from sglang.kernels.ops.attention.pad import ( pad_sequence_with_mask_kernel as pad_sequence_with_mask_kernel, ) -from sglang.srt.layers.attention.triton_ops.pad import ( +from sglang.kernels.ops.attention.pad import ( seqlens_expand_kernel as seqlens_expand_kernel, ) -from sglang.srt.layers.attention.triton_ops.pad import ( +from sglang.kernels.ops.attention.pad import ( seqlens_expand_triton as seqlens_expand_triton, ) -from sglang.srt.layers.attention.triton_ops.rope_cache import ( +from sglang.kernels.ops.kvcache.cache_ops import ( + concat_and_cast_mha_k_kernel as concat_and_cast_mha_k_kernel, +) +from sglang.kernels.ops.kvcache.cache_ops import ( + concat_and_cast_mha_k_triton as concat_and_cast_mha_k_triton, +) +from sglang.kernels.ops.kvcache.cache_ops import ( + launch_reshape_and_cache_flash as launch_reshape_and_cache_flash, +) +from sglang.kernels.ops.kvcache.cache_ops import ( + reshape_and_cache_flash as reshape_and_cache_flash, +) +from sglang.kernels.ops.kvcache.kv_indices import ( + create_flashinfer_kv_indices_triton as create_flashinfer_kv_indices_triton, +) +from sglang.kernels.ops.kvcache.kv_indices import ( + create_flashmla_kv_indices_triton as create_flashmla_kv_indices_triton, +) +from sglang.kernels.ops.kvcache.kv_indices import ( + get_num_kv_index_blocks_flashmla as get_num_kv_index_blocks_flashmla, +) +from sglang.kernels.ops.kvcache.kv_indices import ( + get_num_page_per_block_flashmla as get_num_page_per_block_flashmla, +) +from sglang.kernels.ops.kvcache.rope_cache import ( fused_qk_rope_reshape_and_cache as fused_qk_rope_reshape_and_cache, ) from sglang.srt.utils import is_cuda diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 9b063b015..526a85aa4 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -61,14 +61,14 @@ if _is_npu: if _is_xpu: from sgl_kernel.flash_attn import flash_attn_varlen_func +from sglang.kernels.ops.attention.prefill_attention import ( + context_attention_fwd, +) from sglang.srt.distributed import ( split_tensor_along_last_dim, tensor_model_parallel_all_gather, ) from sglang.srt.distributed import utils as dist_utils -from sglang.srt.layers.attention.triton_ops.prefill_attention import ( - context_attention_fwd, -) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index 68aa893d2..6a315d8a5 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/python/sglang/srt/layers/attention/wave_backend.py @@ -7,11 +7,11 @@ from typing import TYPE_CHECKING, Optional import torch import triton -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.triton_ops.kv_indices import ( +from sglang.kernels.ops.attention.metadata import get_num_kv_splits_triton +from sglang.kernels.ops.kvcache.kv_indices import ( create_flashinfer_kv_indices_triton, ) -from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import get_bool_env_var, get_device_core_count diff --git a/python/sglang/srt/layers/elementwise.py b/python/sglang/srt/layers/elementwise.py index 4065e9dea..1414e0038 100644 --- a/python/sglang/srt/layers/elementwise.py +++ b/python/sglang/srt/layers/elementwise.py @@ -5,7 +5,7 @@ import triton import triton.language as tl from sglang.jit_kernel.utils import is_arch_support_pdl -from sglang.srt.layers.triton_ops.softcap import softcap_out as fused_softcap +from sglang.kernels.ops.activation.softcap import softcap_out as fused_softcap from sglang.srt.utils import is_hip from sglang.srt.utils.custom_op import register_custom_op diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index aad02844a..391cbbb25 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -21,6 +21,9 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch from torch import nn +from sglang.kernels.ops.activation.softcap import ( + softcap_inplace_logits as fused_softcap, +) from sglang.srt.distributed.device_communicators import triton_symm_mem_ag from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( @@ -33,7 +36,6 @@ from sglang.srt.layers.dp_attention import ( get_dp_dtype, get_dp_hidden_size, ) -from sglang.srt.layers.triton_ops.softcap import softcap_inplace_logits as fused_softcap from sglang.srt.layers.utils.logprob import ( InputLogprobsResult, get_token_ids_logprobs_chunk, diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py index 5d6f45666..472bc55a5 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py @@ -16,7 +16,7 @@ _is_xpu = is_xpu() _is_musa = is_musa() if _is_cuda or _is_hip or _is_xpu or _is_musa: - from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size + from sglang.kernels.ops.moe import moe_align_block_size as sgl_moe_align_block_size def moe_align_block_size( diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 7a3c98b5a..007466b52 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -179,7 +179,7 @@ if _is_cuda: fused_topk_deepseek = None if _is_cuda or _is_hip or _is_xpu: - from sgl_kernel import topk_softmax + from sglang.kernels.ops.moe import topk_softmax try: from sgl_kernel import topk_sigmoid diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 0bd72ccb0..0f2b7eef9 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -55,11 +55,10 @@ _is_sm120_supported = is_sm120_supported() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda or _is_musa: - from sgl_kernel import sgl_per_token_quant_fp8 - from sglang.jit_kernel.per_tensor_quant_fp8 import ( per_tensor_quant_fp8 as sgl_per_tensor_quant_fp8, ) + from sglang.kernels.ops.quantization import sgl_per_token_quant_fp8 # Temporary try: diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 7a47d10cd..b6cd3d925 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -3,12 +3,12 @@ from typing import List, Optional, Tuple import torch -from sglang.srt.lora.backend.base_backend import BaseLoRABackend -from sglang.srt.lora.triton_ops import ( +from sglang.kernels.ops.gemm.chunked_embedding_lora_a import ( chunked_embedding_lora_a_forward, - chunked_sgmv_lora_expand_forward, - chunked_sgmv_lora_shrink_forward, ) +from sglang.kernels.ops.gemm.chunked_sgmv_expand import chunked_sgmv_lora_expand_forward +from sglang.kernels.ops.gemm.chunked_sgmv_shrink import chunked_sgmv_lora_shrink_forward +from sglang.srt.lora.backend.base_backend import BaseLoRABackend from sglang.srt.lora.utils import ( LoRABatchInfo, generate_sequence_lengths, diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index 47bca9c6d..e9708f97d 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -3,14 +3,12 @@ from typing import List, Optional, Tuple import torch +from sglang.kernels.ops.gemm.embedding_lora_a import embedding_lora_a_fwd +from sglang.kernels.ops.gemm.gate_up_lora_b import gate_up_lora_b_fwd +from sglang.kernels.ops.gemm.qkv_lora_b import qkv_lora_b_fwd +from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd +from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd from sglang.srt.lora.backend.base_backend import BaseLoRABackend -from sglang.srt.lora.triton_ops import ( - embedding_lora_a_fwd, - gate_up_lora_b_fwd, - qkv_lora_b_fwd, - sgemm_lora_a_fwd, - sgemm_lora_b_fwd, -) from sglang.srt.lora.utils import ( LoRABatchInfo, get_lm_head_pruned_lens, diff --git a/python/sglang/srt/lora/deepseek_mla_correction.py b/python/sglang/srt/lora/deepseek_mla_correction.py index f81d5bcbf..2edc33416 100644 --- a/python/sglang/srt/lora/deepseek_mla_correction.py +++ b/python/sglang/srt/lora/deepseek_mla_correction.py @@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Optional, Tuple import torch -from sglang.srt.lora.triton_ops import ( +from sglang.kernels.ops.gemm.kv_b_lora_absorbed import ( step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, diff --git a/python/sglang/srt/lora/lora_moe_runner_marlin.py b/python/sglang/srt/lora/lora_moe_runner_marlin.py index 3f7540f1f..447846f44 100644 --- a/python/sglang/srt/lora/lora_moe_runner_marlin.py +++ b/python/sglang/srt/lora/lora_moe_runner_marlin.py @@ -23,9 +23,8 @@ if TYPE_CHECKING: _is_cuda = is_cuda() if _is_cuda: - from sgl_kernel import silu_and_mul - from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm + from sglang.kernels.ops.activation import silu_and_mul from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import ( get_scalar_type, ) diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index f44bc486a..927e42985 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -316,10 +316,8 @@ def _add_lora_gate_up_delta( routing_cache: dict | None = None, ) -> None: """Add LoRA gate_up delta to intermediate_cache in-place.""" - from sglang.srt.lora.triton_ops import ( - fused_moe_lora, - merged_experts_fused_moe_lora_add, - ) + from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora + from sglang.kernels.ops.moe.virtual_experts import merged_experts_fused_moe_lora_add if lora_info is None or lora_info.max_lora_rank == 0: return @@ -409,10 +407,8 @@ def _add_lora_down_delta( routing_cache: dict | None = None, ) -> None: """Add LoRA down delta to intermediate_cache in-place.""" - from sglang.srt.lora.triton_ops import ( - fused_moe_lora, - merged_experts_fused_moe_lora_add, - ) + from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora + from sglang.kernels.ops.moe.virtual_experts import merged_experts_fused_moe_lora_add if lora_info.max_lora_rank == 0: return diff --git a/python/sglang/srt/lora/triton_ops/__init__.py b/python/sglang/srt/lora/triton_ops/__init__.py deleted file mode 100644 index c43e73697..000000000 --- a/python/sglang/srt/lora/triton_ops/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -from .chunked_embedding_lora_a import chunked_embedding_lora_a_forward -from .chunked_sgmv_expand import chunked_sgmv_lora_expand_forward -from .chunked_sgmv_shrink import chunked_sgmv_lora_shrink_forward -from .embedding_lora_a import embedding_lora_a_fwd -from .fused_moe_lora_kernel import fused_moe_lora -from .gate_up_lora_b import gate_up_lora_b_fwd -from .kv_b_lora_absorbed import ( - step_a_q_fwd, - step_a_v_fwd, - step_b_q_fwd, - step_b_v_fwd, -) -from .qkv_lora_b import qkv_lora_b_fwd -from .sgemm_lora_a import sgemm_lora_a_fwd -from .sgemm_lora_b import sgemm_lora_b_fwd -from .virtual_experts import merged_experts_fused_moe_lora_add - -__all__ = [ - "gate_up_lora_b_fwd", - "qkv_lora_b_fwd", - "sgemm_lora_a_fwd", - "sgemm_lora_b_fwd", - "chunked_sgmv_lora_shrink_forward", - "chunked_sgmv_lora_expand_forward", - "fused_moe_lora", - "chunked_embedding_lora_a_forward", - "embedding_lora_a_fwd", - "merged_experts_fused_moe_lora_add", - "step_a_q_fwd", - "step_a_v_fwd", - "step_b_q_fwd", - "step_b_v_fwd", -] diff --git a/python/sglang/srt/lora/trtllm_lora_temp/attention.py b/python/sglang/srt/lora/trtllm_lora_temp/attention.py index 410d759e8..cdbaaca83 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/attention.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/attention.py @@ -36,10 +36,8 @@ def qkv_proj_lora_forward(self, input_: torch.Tensor): if not self.set_lora or not is_two_stream_active(input_): return get_original_qkv_forward()(self, input_) - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( - qkv_lora_b_fwd, - sgemm_lora_a_fwd, - ) + from sglang.kernels.ops.gemm.trtllm_lora_temp.qkv_lora_b import qkv_lora_b_fwd + from sglang.kernels.ops.gemm.trtllm_lora_temp.sgemm_lora_a import sgemm_lora_a_fwd bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None side_stream = get_lora_side_stream() @@ -238,7 +236,7 @@ def replicated_lora_forward(self, x: torch.Tensor): base_output=output, ) else: - from sglang.srt.lora.trtllm_lora_temp.triton_ops import qkv_lora_b_fwd + from sglang.kernels.ops.gemm.trtllm_lora_temp.qkv_lora_b import qkv_lora_b_fwd output = qkv_lora_b_fwd( lora_a_output, diff --git a/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py b/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py index b2e066e3a..7ae19430d 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py @@ -29,10 +29,18 @@ step_a_q_fwd = step_a_v_fwd = step_b_q_fwd = step_b_v_fwd = None def _ensure_step_kernels() -> None: global step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, step_b_v_fwd if step_a_q_fwd is None: - from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_a_q_fwd as _aq - from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_a_v_fwd as _av - from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_b_q_fwd as _bq - from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_b_v_fwd as _bv + from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import ( + step_a_q_fwd as _aq, + ) + from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import ( + step_a_v_fwd as _av, + ) + from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import ( + step_b_q_fwd as _bq, + ) + from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import ( + step_b_v_fwd as _bv, + ) step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, step_b_v_fwd = _aq, _av, _bq, _bv diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py index cebed4c06..a361d9af4 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py @@ -52,6 +52,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora( trtllm_fp8_block_scale_routed_moe_lora, ) from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk + from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( + merged_experts_fused_moe_lora_add, + ) from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType @@ -62,9 +65,6 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora( from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import ( maybe_overlap_staged_shared_add, ) - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( - merged_experts_fused_moe_lora_add, - ) from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode assert runner_config.activation == "silu" and runner_config.is_gated, ( @@ -316,6 +316,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora( """ from sglang.jit_kernel.trtllm_lora_temp import trtllm_bf16_routed_moe_lora from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk + from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( + merged_experts_fused_moe_lora_add, + ) from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( fused_experts_none_to_flashinfer_trtllm_bf16, get_activation_type, @@ -323,9 +326,6 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora( from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( - merged_experts_fused_moe_lora_add, - ) from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode assert ( @@ -469,14 +469,14 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora( trtllm_fp4_block_scale_routed_moe_lora, ) from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk + from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( + merged_experts_fused_moe_lora_add, + ) from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( fused_experts_none_to_flashinfer_trtllm_fp4, ) from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( - merged_experts_fused_moe_lora_add, - ) from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode assert ( diff --git a/python/sglang/srt/lora/trtllm_lora_temp/merged_column.py b/python/sglang/srt/lora/trtllm_lora_temp/merged_column.py index 88b322c93..c849cea89 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/merged_column.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/merged_column.py @@ -32,11 +32,11 @@ def merged_column_lora_forward(self, input_: torch.Tensor): if not self.set_lora or not is_two_stream_active(input_): return get_original_merged_column_forward()(self, input_) - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( + from sglang.kernels.ops.gemm.trtllm_lora_temp.gate_up_lora_b import ( gate_up_lora_b_fwd, - qkv_lora_b_fwd, - sgemm_lora_a_fwd, ) + from sglang.kernels.ops.gemm.trtllm_lora_temp.qkv_lora_b import qkv_lora_b_fwd + from sglang.kernels.ops.gemm.trtllm_lora_temp.sgemm_lora_a import sgemm_lora_a_fwd bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None side_stream = get_lora_side_stream() diff --git a/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py b/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py index e0699c60f..b21ff0378 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py @@ -58,6 +58,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream( trtllm_fp8_block_scale_routed_moe_lora, ) from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk + from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( + merged_experts_fused_moe_lora_add, + ) from sglang.srt.distributed import get_tp_group from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -70,9 +73,6 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream( from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import ( maybe_overlap_staged_shared_add, ) - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( - merged_experts_fused_moe_lora_add, - ) from sglang.srt.utils.common import next_power_of_2 assert runner_config.activation == "silu" and runner_config.is_gated, ( @@ -362,6 +362,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream( trtllm_fp4_block_scale_routed_moe_lora, ) from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk + from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( + merged_experts_fused_moe_lora_add, + ) from sglang.srt.distributed import get_tp_group from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -369,9 +372,6 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream( from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( - merged_experts_fused_moe_lora_add, - ) assert ( runner_config.activation == "silu" and runner_config.is_gated @@ -608,6 +608,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream( # ---- two-stream fast path ---- from sglang.jit_kernel.trtllm_lora_temp import trtllm_bf16_routed_moe_lora from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk + from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import ( + merged_experts_fused_moe_lora_add, + ) from sglang.srt.distributed import get_tp_group from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -619,9 +622,6 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream( from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType - from sglang.srt.lora.trtllm_lora_temp.triton_ops import ( - merged_experts_fused_moe_lora_add, - ) assert ( runner_config.activation == "silu" and runner_config.is_gated diff --git a/python/sglang/srt/lora/trtllm_lora_temp/specialized_expand.py b/python/sglang/srt/lora/trtllm_lora_temp/specialized_expand.py index 6461a473a..e20a1d82a 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/specialized_expand.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/specialized_expand.py @@ -6,7 +6,7 @@ It is rank-specialized: the ``R`` dimension (LoRA rank) is a triton specialization (R=16, R=32, R=64 are all supported up to the ``R <= 64`` assert, with no perf interaction between them — each gets its own kernel). -Called from :mod:`sglang.srt.lora.triton_ops.virtual_experts` when +Called from :mod:`sglang.kernels.ops.moe.virtual_experts` when ``use_direct_expand_add=True`` (the trtllm-lora path uses this when ``max_lora_rank <= 64``); the generic ``invoke_fused_moe_kernel`` is used when that flag is False (incl. ranks above 64). diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/__init__.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/__init__.py deleted file mode 100644 index 6b60a3aea..000000000 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/__init__.py +++ /dev/null @@ -1,32 +0,0 @@ -"""Experimental TRT-LLM LoRA kernels (copies of the upstream triton_ops kernels -with the SGLANG_EXPERIMENTAL_LORA_OPTI optimizations). - -These are forked from ``sglang.srt.lora.triton_ops`` so the upstream kernels stay -byte-pristine; only the experimental forwards / dispatch in this package import -from here. The opt branches inside each kernel are still gated by ``lora_envs`` -(master-gated by SGLANG_EXPERIMENTAL_LORA_OPTI). -""" - -from .gate_up_lora_b import gate_up_lora_b_fwd -from .kv_b_lora_absorbed import ( - step_a_q_fwd, - step_a_v_fwd, - step_b_q_fwd, - step_b_v_fwd, -) -from .qkv_lora_b import qkv_lora_b_fwd -from .sgemm_lora_a import sgemm_lora_a_fwd -from .sgemm_lora_b import sgemm_lora_b_fwd -from .virtual_experts import merged_experts_fused_moe_lora_add - -__all__ = [ - "gate_up_lora_b_fwd", - "qkv_lora_b_fwd", - "sgemm_lora_a_fwd", - "sgemm_lora_b_fwd", - "merged_experts_fused_moe_lora_add", - "step_a_q_fwd", - "step_a_v_fwd", - "step_b_q_fwd", - "step_b_v_fwd", -] diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 6e6e89e80..d27c89ad6 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -5,8 +5,8 @@ from typing import TYPE_CHECKING, Optional, Sequence import torch +from sglang.kernels.ops.speculative.gather_spec_extras import gather_spec_extras from sglang.srt.environ import envs -from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras from sglang.srt.utils import is_cuda, is_hip, is_npu if TYPE_CHECKING: diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index 2531396c7..eb34cf7a5 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -24,11 +24,11 @@ from typing import TYPE_CHECKING import torch -from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator -from sglang.srt.mem_cache.triton_ops.allocator import ( +from sglang.kernels.ops.memory.allocator import ( alloc_decode_kernel, alloc_extend_kernel, ) +from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator from sglang.srt.utils import ( get_bool_env_var, get_num_new_pages, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index a58f8970b..e2cedf6ae 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -6,6 +6,15 @@ from typing import TYPE_CHECKING, Optional import numpy as np import torch +from sglang.kernels.ops.memory.common import ( + _get_last_loc_safe_kernel as _get_last_loc_safe_kernel, +) +from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel +from sglang.kernels.ops.memory.common import ( + get_last_loc_triton, + get_last_loc_triton_safe, + write_req_to_token_pool_triton, +) from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_evict_dsv4_state_on_swa, maybe_write_dsv4_decode, @@ -14,17 +23,6 @@ from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool -from sglang.srt.mem_cache.triton_ops.common import ( - _get_last_loc_safe_kernel as _get_last_loc_safe_kernel, -) -from sglang.srt.mem_cache.triton_ops.common import ( - get_last_loc_kernel as get_last_loc_kernel, -) -from sglang.srt.mem_cache.triton_ops.common import ( - get_last_loc_triton, - get_last_loc_triton_safe, - write_req_to_token_pool_triton, -) from sglang.srt.runtime_context import get_server_args from sglang.srt.server_args import ServerArgs from sglang.srt.utils import is_cuda, is_hip, is_npu, support_triton diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index b2856d0d7..80c5248b9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -36,6 +36,11 @@ import triton import triton.language as tl from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache +from sglang.kernels.ops.kvcache.cache_move import ( + copy_all_layer_kv_cache_func, + set_kv_buffer_prefix_valid_tiled, + store_cache_4d, +) from sglang.srt.configs.mamba_utils import BaseLinearStateParams from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs @@ -60,11 +65,6 @@ from sglang.srt.mem_cache.layout.page_major import ( mamba_entry_bytes, mha_entry_bytes, ) -from sglang.srt.mem_cache.triton_ops.cache_move import ( - copy_all_layer_kv_cache_func, - set_kv_buffer_prefix_valid_tiled, - store_cache_4d, -) from sglang.srt.mem_cache.utils import ( get_mla_kv_buffer_triton, maybe_init_custom_mem_pool, diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 921badbaf..179eebdb0 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -30,6 +30,7 @@ from typing import Dict, List, Optional, Set, Tuple import torch from torch.profiler import record_function +from sglang.kernels.ops.memory.virtual_slot import alloc_bind_inplace from sglang.srt.environ import envs from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.paged import ( @@ -37,7 +38,6 @@ from sglang.srt.mem_cache.allocator.paged import ( alloc_extend_kernel, ) from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.triton_ops.virtual_slot import alloc_bind_inplace from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool from sglang.srt.utils.common import get_num_new_pages, next_power_of_2 diff --git a/python/sglang/srt/mem_cache/triton_ops/__init__.py b/python/sglang/srt/mem_cache/triton_ops/__init__.py deleted file mode 100644 index 3e69762f9..000000000 --- a/python/sglang/srt/mem_cache/triton_ops/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Triton kernels for memory cache operations.""" diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index bb02e9096..bef6ff73d 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -32,6 +32,7 @@ import torch import triton from torch.profiler import record_function +from sglang.kernels.ops.kvcache.cache_move import store_cache_4d_kernel from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.mem_cache.layout.page_major import ( build_page_major_mamba_views, @@ -45,7 +46,6 @@ from sglang.srt.mem_cache.memory_pool import ( unwrap_write_loc, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool -from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d_kernel from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 85d30fbd7..9cb4b1623 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -15,6 +15,30 @@ from typing import Any, Callable, List, Optional, Tuple +from sglang.kernels.ops.kvcache.mla_buffer import ( + get_mla_kv_buffer_kernel as get_mla_kv_buffer_kernel, +) +from sglang.kernels.ops.kvcache.mla_buffer import ( + get_mla_kv_buffer_triton as get_mla_kv_buffer_triton, +) +from sglang.kernels.ops.kvcache.mla_buffer import ( + set_mla_kv_buffer_fp8_quant_kernel as set_mla_kv_buffer_fp8_quant_kernel, +) +from sglang.kernels.ops.kvcache.mla_buffer import ( + set_mla_kv_buffer_kernel as set_mla_kv_buffer_kernel, +) +from sglang.kernels.ops.kvcache.mla_buffer import ( + set_mla_kv_buffer_triton as set_mla_kv_buffer_triton, +) +from sglang.kernels.ops.kvcache.mla_buffer import ( + set_mla_kv_buffer_triton_fp8_quant as set_mla_kv_buffer_triton_fp8_quant, +) +from sglang.kernels.ops.kvcache.mla_buffer import ( + set_mla_kv_scale_buffer_kernel as set_mla_kv_scale_buffer_kernel, +) +from sglang.kernels.ops.kvcache.mla_buffer import ( + set_mla_kv_scale_buffer_triton as set_mla_kv_scale_buffer_triton, +) from sglang.srt.environ import envs from sglang.srt.mem_cache.cpp_utils.native_hash import get_native_hash from sglang.srt.mem_cache.evict_policy import ( @@ -27,30 +51,6 @@ from sglang.srt.mem_cache.evict_policy import ( PriorityStrategy, SLRUStrategy, ) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - get_mla_kv_buffer_kernel as get_mla_kv_buffer_kernel, -) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - get_mla_kv_buffer_triton as get_mla_kv_buffer_triton, -) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - set_mla_kv_buffer_fp8_quant_kernel as set_mla_kv_buffer_fp8_quant_kernel, -) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - set_mla_kv_buffer_kernel as set_mla_kv_buffer_kernel, -) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - set_mla_kv_buffer_triton as set_mla_kv_buffer_triton, -) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - set_mla_kv_buffer_triton_fp8_quant as set_mla_kv_buffer_triton_fp8_quant, -) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - set_mla_kv_scale_buffer_kernel as set_mla_kv_scale_buffer_kernel, -) -from sglang.srt.mem_cache.triton_ops.mla_buffer import ( - set_mla_kv_scale_buffer_triton as set_mla_kv_scale_buffer_triton, -) _EVICTION_POLICY_FACTORIES: dict[str, Callable[[], EvictionStrategy]] = { "lru": LRUStrategy, diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py index b8b1e12bf..770958c76 100644 --- a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py +++ b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py @@ -5,11 +5,11 @@ from typing import List, Optional import torch -from sglang.srt.environ import envs -from sglang.srt.layers.attention.triton_ops.kv_indices import ( +from sglang.kernels.ops.kvcache.kv_indices import ( create_chunked_prefix_cache_kv_indices, create_flashinfer_kv_indices_triton, ) +from sglang.srt.environ import envs from sglang.srt.model_executor.forward_context import ( get_req_to_token_pool, get_token_to_kv_pool, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index e1fb9b28e..9b29865b2 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -36,6 +36,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union import torch +from sglang.kernels.ops.attention.position import compute_position_triton from sglang.srt.environ import envs from sglang.srt.kv_canary.req_to_expected_token_ids_manager import ( compute_req_all_ids_info, @@ -48,7 +49,6 @@ from sglang.srt.layers.dp_attention import ( from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( ForwardBatchDeepSeekMHAMixin, ) -from sglang.srt.model_executor.triton_ops.position import compute_position_triton from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.utils import ( is_cuda, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py index 65545069e..f46e49c1d 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py @@ -24,7 +24,7 @@ if _is_cuda: from sgl_kernel import bmm_fp8 if _is_hip: - from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import ( + from sglang.kernels.ops.attention.rocm_mla_decode_rope import ( decode_attention_fwd_grouped_rope, ) diff --git a/python/sglang/srt/models/triton_ops/deepseek_v4.py b/python/sglang/srt/models/triton_ops/deepseek_v4.py deleted file mode 100644 index 62915b128..000000000 --- a/python/sglang/srt/models/triton_ops/deepseek_v4.py +++ /dev/null @@ -1,56 +0,0 @@ -from typing import Optional - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _rms_normalize_kernel( - x_ptr, - weight_ptr, - eps, - stride_row, - dim, - BLOCK_SIZE: tl.constexpr, - HAS_WEIGHT: tl.constexpr, -): - pid = tl.program_id(0) - - offs = tl.arange(0, BLOCK_SIZE) - mask = offs < dim - - base = pid * stride_row - x = tl.load(x_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - - mean_sq = tl.sum(x * x, axis=0) / dim - rms_inv = tl.rsqrt(mean_sq + eps) - out = x * rms_inv - - if HAS_WEIGHT: - weight = tl.load(weight_ptr + offs, mask=mask, other=0.0) - out = out * weight - - tl.store(x_ptr + base + offs, out, mask=mask) - - -def rms_normalize_triton( - x: torch.Tensor, eps: float, weight: Optional[torch.Tensor] = None -) -> torch.Tensor: - dim = x.shape[-1] - x_flat = x.view(-1, dim) - num_rows = x_flat.shape[0] - - BLOCK_SIZE = triton.next_power_of_2(dim) - grid = (num_rows,) - - _rms_normalize_kernel[grid]( - x_flat, - weight, - eps, - x_flat.stride(0), - dim, - BLOCK_SIZE=BLOCK_SIZE, - HAS_WEIGHT=(weight is not None), - ) - return x diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index 909bdbdf1..c871798be 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -190,13 +190,13 @@ class EagleDraftWorkerBase(ABC): topk: int, num_steps: int, ): + from sglang.kernels.ops.speculative.cache_locs import ( + assign_draft_cache_locs_contiguous, + ) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ) - from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_draft_cache_locs_contiguous, - ) if not batch.forward_mode.is_idle(): bs = len(batch.seq_lens) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index a35363658..68c5ae526 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -4,6 +4,11 @@ from typing import List, Optional import torch +from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func +from sglang.kernels.ops.speculative.dflash import ( + _compute_dflash_accept_bonus_triton_unchecked, + _prepare_dflash_draft_block_unchecked, +) from sglang.srt.distributed import get_tp_group from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult @@ -29,11 +34,6 @@ from sglang.srt.speculative.dflash_utils import ( ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func -from sglang.srt.speculative.triton_ops.cache_locs import assign_extend_cache_locs_func -from sglang.srt.speculative.triton_ops.dflash import ( - _compute_dflash_accept_bonus_triton_unchecked, - _prepare_dflash_draft_block_unchecked, -) from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu _is_npu = is_npu() @@ -47,7 +47,7 @@ _FusedKVMaterializeHelper = None def _get_fused_kv_materialize_helper(): global _FusedKVMaterializeHelper if _FusedKVMaterializeHelper is None: - from sglang.srt.speculative.triton_ops.fused_kv_materialize import ( + from sglang.kernels.ops.speculative.fused_kv_materialize import ( FusedKVMaterializeHelper, ) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index b397a6b1a..277b15831 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -8,6 +8,10 @@ from typing import TYPE_CHECKING, List, Optional import torch +from sglang.kernels.ops.speculative.spec_tree import ( + sgl_build_tree_kernel_efficient_triton, + verify_tree_greedy_kernel_triton, +) from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import ( alloc_paged_token_slots_extend_npu, ) @@ -21,10 +25,6 @@ from sglang.srt.mem_cache.common import ( get_last_loc, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.speculative.triton_ops.spec_tree import ( - sgl_build_tree_kernel_efficient_triton, - verify_tree_greedy_kernel_triton, -) from sglang.srt.utils import ( is_cpu, is_cuda, @@ -504,15 +504,15 @@ def eagle_prepare_for_verify( batch: ScheduleBatch, target_worker: TpModelWorker, ): + from sglang.kernels.ops.speculative.cache_locs import ( + assign_extend_cache_locs_func, + ) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ForwardMode, ) from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify - from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_extend_cache_locs_func, - ) if not batch.forward_mode.is_idle(): # Assign cache locations diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 624dedf1b..80526c3f5 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -5,6 +5,7 @@ from typing import List, Optional, Tuple import torch +from sglang.kernels.ops.speculative.eagle import fill_bonus_tokens_func from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import ( EAGLEDraftExtendNpuGraphRunner, @@ -90,7 +91,6 @@ from sglang.srt.speculative.spec_utils import ( select_top_k_tokens, spec_stage_span, ) -from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens_func from sglang.srt.utils.async_probe import ( maybe_detect_inf, maybe_detect_nan, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_utils.py b/python/sglang/srt/speculative/multi_layer_eagle_utils.py index bc54df1ca..f7f754d8c 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_utils.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_utils.py @@ -12,7 +12,7 @@ # limitations under the License. # ============================================================================== -from sglang.srt.speculative.triton_ops.multi_layer_eagle import ( +from sglang.kernels.ops.speculative.multi_layer_eagle import ( rotate_input_ids, rotate_input_ids_kernel, ) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 873848a46..de5d42685 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple import torch +from sglang.kernels.ops.speculative.eagle import fill_bonus_tokens_func from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.graph_runner.multi_layer_eagle_draft_extend_npu_graph_runner import ( MultiLayerEagleMultiStepDraftExtendNpuGraphRunner, @@ -67,7 +68,6 @@ from sglang.srt.speculative.spec_utils import ( sample_draft_proposal, select_top_k_tokens, ) -from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens_func from sglang.srt.utils import is_cpu, is_npu from sglang.srt.utils.async_probe import ( maybe_detect_inf, diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 90524c74c..ed1d657fb 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -5,6 +5,9 @@ import numpy as np import torch from sgl_kernel.speculative import reconstruct_indices_from_tree_mask +from sglang.kernels.ops.speculative.cache_locs import ( + assign_extend_cache_locs_func as assign_extend_cache_locs_func, +) from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult @@ -23,9 +26,6 @@ from sglang.srt.speculative.spec_utils import ( prepare_mamba_track_for_verify, record_stream_for_v2_verify, ) -from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_extend_cache_locs_func as assign_extend_cache_locs_func, -) from sglang.srt.utils import is_cpu from sglang.srt.utils.async_probe import maybe_detect_inf, maybe_detect_nan diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 8d75142ab..437ea74b3 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -9,6 +9,33 @@ from typing import TYPE_CHECKING, List, Optional import torch from huggingface_hub import snapshot_download +from sglang.kernels.ops.speculative.cache_locs import ( + align_evict_mask_to_page_size as align_evict_mask_to_page_size, +) +from sglang.kernels.ops.speculative.cache_locs import ( + assign_extend_cache_locs as assign_extend_cache_locs, +) +from sglang.kernels.ops.speculative.cache_locs import ( + assign_req_to_token_pool as assign_req_to_token_pool, +) +from sglang.kernels.ops.speculative.cache_locs import ( + assign_req_to_token_pool_func as assign_req_to_token_pool_func, +) +from sglang.kernels.ops.speculative.cache_locs import ( + filter_finished_cache_loc_kernel as filter_finished_cache_loc_kernel, +) +from sglang.kernels.ops.speculative.cache_locs import ( + generate_draft_decode_kv_indices as generate_draft_decode_kv_indices, +) +from sglang.kernels.ops.speculative.cache_locs import ( + get_src_tgt_cache_loc as get_src_tgt_cache_loc, +) +from sglang.kernels.ops.speculative.cache_locs import ( + get_target_cache_loc as get_target_cache_loc, +) +from sglang.kernels.ops.speculative.eagle import ( + fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func, +) from sglang.srt.distributed.parallel_state import ( GroupCoordinator, patch_tensor_parallel_group, @@ -16,33 +43,6 @@ from sglang.srt.distributed.parallel_state import ( from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import set_mamba_track_indices_from_reqs from sglang.srt.runtime_context import get_server_args -from sglang.srt.speculative.triton_ops.cache_locs import ( - align_evict_mask_to_page_size as align_evict_mask_to_page_size, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_extend_cache_locs as assign_extend_cache_locs, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_req_to_token_pool as assign_req_to_token_pool, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_req_to_token_pool_func as assign_req_to_token_pool_func, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - filter_finished_cache_loc_kernel as filter_finished_cache_loc_kernel, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - generate_draft_decode_kv_indices as generate_draft_decode_kv_indices, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - get_src_tgt_cache_loc as get_src_tgt_cache_loc, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - get_target_cache_loc as get_target_cache_loc, -) -from sglang.srt.speculative.triton_ops.eagle import ( - fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func, -) from sglang.srt.utils import ( is_cpu, is_cuda, diff --git a/python/sglang/srt/speculative/triton_ops/__init__.py b/python/sglang/srt/speculative/triton_ops/__init__.py deleted file mode 100644 index a8ea8f4c7..000000000 --- a/python/sglang/srt/speculative/triton_ops/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Triton kernels for speculative decoding.""" - -from sglang.srt.speculative.triton_ops.fused_kv_materialize import ( - FusedKVMaterializeHelper, -) - -__all__ = ["FusedKVMaterializeHelper"] diff --git a/test/manual/lora/test_lora_tuning_config.py b/test/manual/lora/test_lora_tuning_config.py index 6601a7bde..002c97cec 100644 --- a/test/manual/lora/test_lora_tuning_config.py +++ b/test/manual/lora/test_lora_tuning_config.py @@ -6,7 +6,7 @@ import tempfile import unittest from unittest.mock import patch -from sglang.srt.lora.triton_ops.lora_tuning_config import ( +from sglang.kernels.ops.gemm.lora_tuning_config import ( DEFAULT_EXPAND_CONFIG, DEFAULT_SHRINK_CONFIG, get_lora_config_file_name, @@ -15,7 +15,7 @@ from sglang.srt.lora.triton_ops.lora_tuning_config import ( get_lora_shrink_config, ) -_MODULE = "sglang.srt.lora.triton_ops.lora_tuning_config" +_MODULE = "sglang.kernels.ops.gemm.lora_tuning_config" # Shared fixture _TUNED_CONFIGS = { @@ -88,7 +88,7 @@ class TestConfigSelection(unittest.TestCase): def setUp(self): get_lora_configs.cache_clear() - from sglang.srt.lora.triton_ops import lora_tuning_config + from sglang.kernels.ops.gemm import lora_tuning_config lora_tuning_config._logged_configs.clear() diff --git a/test/manual/test_triton_attention_rocm_mla.py b/test/manual/test_triton_attention_rocm_mla.py index f3074aeeb..7c89d521b 100644 --- a/test/manual/test_triton_attention_rocm_mla.py +++ b/test/manual/test_triton_attention_rocm_mla.py @@ -3,10 +3,10 @@ import unittest import torch -from sglang.srt.layers.attention.triton_ops.decode_attention import ( +from sglang.kernels.ops.attention.decode_attention import ( decode_attention_fwd_grouped, ) -from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import ( +from sglang.kernels.ops.attention.rocm_mla_decode_rope import ( decode_attention_fwd_grouped_rope, ) from sglang.srt.layers.rotary_embedding import DeepseekScalingRotaryEmbedding diff --git a/test/manual/test_trtllm_fp8_kv_kernel.py b/test/manual/test_trtllm_fp8_kv_kernel.py index c97deaf11..1ff63d781 100644 --- a/test/manual/test_trtllm_fp8_kv_kernel.py +++ b/test/manual/test_trtllm_fp8_kv_kernel.py @@ -6,7 +6,7 @@ import unittest import torch -from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import ( +from sglang.kernels.ops.kvcache.trtllm_fp8_kv_kernel import ( fused_fp8_set_kv_buffer, ) from sglang.test.test_utils import CustomTestCase diff --git a/test/registered/attention/test_triton_attention_kernels.py b/test/registered/attention/test_triton_attention_kernels.py index f134cd5c0..10d2e9717 100644 --- a/test/registered/attention/test_triton_attention_kernels.py +++ b/test/registered/attention/test_triton_attention_kernels.py @@ -4,18 +4,18 @@ import unittest import torch import torch.nn.functional as F -from sglang.srt.layers.attention.triton_ops.decode_attention import ( +from sglang.kernels.ops.attention.decode_attention import ( decode_attention_fwd, decode_attention_fwd_grouped, decode_attention_fwd_normal, ) -from sglang.srt.layers.attention.triton_ops.extend_attention import ( +from sglang.kernels.ops.attention.extend_attention import ( build_unified_kv_indices, extend_attention_fwd, extend_attention_fwd_unified, redundant_attention, ) -from sglang.srt.layers.attention.triton_ops.prefill_attention import ( +from sglang.kernels.ops.attention.prefill_attention import ( context_attention_fwd, ) from sglang.srt.utils import get_device @@ -320,7 +320,7 @@ class TestTritonAttention(CustomTestCase): self._test_extend_attention_once(19, 12331, 12, 4, value) def test_extend_attention_block_sizes(self): - from sglang.srt.layers.attention.triton_ops import extend_attention as ea + from sglang.kernels.ops.attention import extend_attention as ea if not ea._is_hip: self.skipTest("HIP-only block-size selection") diff --git a/test/registered/attention/test_trtllm_mha_graph_metadata.py b/test/registered/attention/test_trtllm_mha_graph_metadata.py index eaf9b64fd..a2bca9510 100644 --- a/test/registered/attention/test_trtllm_mha_graph_metadata.py +++ b/test/registered/attention/test_trtllm_mha_graph_metadata.py @@ -12,7 +12,7 @@ import pytest import torch import sglang.srt.layers.attention.trtllm_mha_backend as trtllm_mha_backend -from sglang.srt.layers.attention.triton_ops.trtllm_mha_graph_metadata import ( +from sglang.kernels.ops.kvcache.trtllm_mha_graph_metadata import ( Q_MODE_CUMSUM, Q_MODE_NONE, Q_MODE_STRIDED, diff --git a/test/registered/attention/test_trtllm_mha_page_table.py b/test/registered/attention/test_trtllm_mha_page_table.py index 161720547..32229b43b 100644 --- a/test/registered/attention/test_trtllm_mha_page_table.py +++ b/test/registered/attention/test_trtllm_mha_page_table.py @@ -18,7 +18,7 @@ from typing import Optional import torch -from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import ( +from sglang.kernels.ops.kvcache.trtllm_mha_page_table import ( build_trtllm_mha_page_table, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/attention/test_verify_splitkv.py b/test/registered/attention/test_verify_splitkv.py index 9895e2fae..72123bda2 100644 --- a/test/registered/attention/test_verify_splitkv.py +++ b/test/registered/attention/test_verify_splitkv.py @@ -20,10 +20,10 @@ import unittest import torch -from sglang.srt.layers.attention.triton_ops.extend_attention import ( +from sglang.kernels.ops.attention.extend_attention import ( extend_attention_fwd, ) -from sglang.srt.layers.attention.triton_ops.verify_splitkv import ( +from sglang.kernels.ops.attention.verify_splitkv import ( can_handle, verify_splitkv_fwd, ) diff --git a/test/registered/attention/test_wave_attention_kernels.py b/test/registered/attention/test_wave_attention_kernels.py index 9ffaaaf12..ebbd4f391 100644 --- a/test/registered/attention/test_wave_attention_kernels.py +++ b/test/registered/attention/test_wave_attention_kernels.py @@ -3,14 +3,14 @@ import unittest import torch -from sglang.srt.layers.attention.triton_ops.decode_attention import ( +from sglang.kernels.ops.attention.decode_attention import ( decode_attention_fwd_grouped as triton_decode_attention_fwd_grouped, ) -from sglang.srt.layers.attention.triton_ops.extend_attention import ( +from sglang.kernels.ops.attention.extend_attention import ( extend_attention_fwd, redundant_attention, ) -from sglang.srt.layers.attention.triton_ops.prefill_attention import ( +from sglang.kernels.ops.attention.prefill_attention import ( context_attention_fwd, ) from sglang.srt.layers.attention.wave_ops.decode_attention import ( diff --git a/test/registered/kernels/test_dsa_metadata.py b/test/registered/kernels/test_dsa_metadata.py index de1dc0950..9e62ede04 100644 --- a/test/registered/kernels/test_dsa_metadata.py +++ b/test/registered/kernels/test_dsa_metadata.py @@ -2,7 +2,7 @@ import unittest import torch -from sglang.srt.layers.attention.triton_ops.dsa_metadata import ( +from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_decode_metadata, fused_dsa_draft_extend_metadata, fused_dsa_target_verify_metadata, diff --git a/test/registered/kernels/test_fused_op.py b/test/registered/kernels/test_fused_op.py new file mode 100644 index 000000000..93e0748b1 --- /dev/null +++ b/test/registered/kernels/test_fused_op.py @@ -0,0 +1,290 @@ +"""GPU-free unit tests for ``sglang.kernels``: BaseFusedOp + registry/selector/spec. + +Part of RFC #29630, Phase 2. Covers the multi-backend operator contract +(structural backend detection, priority dispatch, forced backend, runtime +eligibility, tracing), the registry/selector units in isolation, and the +pure-torch reference implementations of the reworked layernorm / activation +ops. Runs in the CPU CI lane; every-backend-vs-native parity lives in +``test_fused_op_gpu_parity.py``. +""" + +import unittest + +import torch + +import sglang.kernels as K +from sglang.kernels.fused_op import BaseFusedOp +from sglang.kernels.registry import KernelRegistry +from sglang.kernels.spec import ( + CapabilityRequirement, + KernelBackend, + KernelSpec, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="base-a-test-cpu") + + +class _ToyAddOp(BaseFusedOp): + """Toy op: element-wise a + b, with a fake 'triton' backend.""" + + op = "test.toy_add" + priority = (KernelBackend.TRITON, KernelBackend.TORCH) + + def forward_native(self, a, b): + return a + b + + def forward_triton(self, a, b): + # Marker so tests can tell which backend ran. + return a + b + 1000 + + +class _CudaOnlyToyOp(BaseFusedOp): + """Toy op whose optimized backend requires CUDA (never eligible on CPU).""" + + op = "test.toy_cuda_only" + priority = (KernelBackend.CUDA_AOT, KernelBackend.TORCH) + capabilities = {KernelBackend.CUDA_AOT: CapabilityRequirement(requires_cuda=True)} + + def forward_native(self, a): + return a * 2 + + def forward_cuda_aot(self, a): + raise AssertionError("must not be selected on a CPU-only box") + + +class TestBaseFusedOp(unittest.TestCase): + def tearDown(self): + K.set_fused_op_backend(None) + K.disable_fused_op_trace() + K.clear_fused_op_trace() + + def test_structural_backend_detection(self): + backends = set(_ToyAddOp().available_backends()) + self.assertEqual( + backends, + {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE, KernelBackend.TRITON}, + ) + + def test_native_always_available(self): + backends = _CudaOnlyToyOp().available_backends() + self.assertIn(KernelBackend.TORCH, backends) + self.assertIn(KernelBackend.TORCH_COMPILE, backends) + + def test_priority_dispatch(self): + op = _ToyAddOp() + a, b = torch.tensor([1.0]), torch.tensor([2.0]) + # TRITON is first in priority and always eligible (no capability). + self.assertEqual(op(a, b).item(), 1003.0) + + def test_explicit_backend_overrides_priority(self): + op = _ToyAddOp() + a, b = torch.tensor([1.0]), torch.tensor([2.0]) + self.assertEqual(op.forward(a, b, backend=KernelBackend.TORCH).item(), 3.0) + + def test_capability_gates_runtime_eligibility(self): + # On a CPU-only box the CUDA backend is filtered out and auto-selection + # falls back to native instead of raising. + op = _CudaOnlyToyOp() + if K.PlatformInfo.detect().is_cuda: + self.skipTest("test requires a CPU-only environment") + self.assertEqual(op(torch.tensor([3.0])).item(), 6.0) + + def test_forced_backend_global_switch(self): + op = _ToyAddOp() + a, b = torch.tensor([1.0]), torch.tensor([2.0]) + K.set_fused_op_backend(KernelBackend.TORCH) + self.assertEqual(op(a, b).item(), 3.0) + K.set_fused_op_backend(None) + self.assertEqual(op(a, b).item(), 1003.0) + + def test_forced_backend_env_var(self): + import sglang.kernels.fused_op as fused_op_module + from sglang.srt.environ import envs + + op = _ToyAddOp() + a, b = torch.tensor([1.0]), torch.tensor([2.0]) + with envs.SGLANG_FORCE_FUSED_OP_BACKEND.override("torch"): + # Reset the module cache so the env var is re-read. + fused_op_module._forced_backend = fused_op_module._UNRESOLVED + self.assertEqual(K.get_fused_op_backend(), KernelBackend.TORCH) + self.assertEqual(op(a, b).item(), 3.0) + fused_op_module._forced_backend = fused_op_module._UNRESOLVED + + def test_unimplemented_backend_raises(self): + op = _ToyAddOp() + with self.assertRaises(NotImplementedError): + op.forward( + torch.tensor([1.0]), + torch.tensor([2.0]), + backend=KernelBackend.CUDA_AOT, + ) + + def test_torch_compile_backend(self): + op = _ToyAddOp() + a, b = torch.tensor([1.0]), torch.tensor([2.0]) + try: + result = op.forward(a, b, backend=KernelBackend.TORCH_COMPILE) + except Exception as e: # inductor toolchain missing in some CI images + self.skipTest(f"torch.compile unavailable: {e}") + self.assertEqual(result.item(), 3.0) + + def test_trace_records_op_backend_and_shapes(self): + op = _ToyAddOp() + K.enable_fused_op_trace() + op(torch.zeros(2, 3), torch.zeros(2, 3)) + records = K.get_fused_op_trace() + self.assertEqual(len(records), 1) + self.assertEqual(records[0].op, "test.toy_add") + self.assertEqual(records[0].backend, "triton") + self.assertEqual( + records[0].tensor_args, + ("torch.float32[2, 3]", "torch.float32[2, 3]"), + ) + + def test_register_fused_op_specs(self): + op = K.registry.get("layernorm.rmsnorm") + backends = {s.backend for s in op} + self.assertEqual( + backends, + { + KernelBackend.TORCH, + KernelBackend.TORCH_COMPILE, + KernelBackend.CUDA_JIT, + KernelBackend.CUDA_AOT, + }, + ) + # Dotted targets resolve to the bound backend methods. + native = K.registry.get_backend("layernorm.rmsnorm", KernelBackend.TORCH) + fn = native.load() + x = torch.randn(4, 64) + w = torch.randn(64) + self.assertTrue(torch.allclose(fn(x, w), _ref_rmsnorm(x, w, 1e-6))) + + +class TestKernelRegistryUnit(unittest.TestCase): + """Isolated KernelRegistry behavior (fresh instance, no global state).""" + + def _spec(self, op="g.n", backend=KernelBackend.TORCH, target="math:sqrt"): + return KernelSpec(op=op, backend=backend, target=target) + + def test_register_and_get(self): + reg = KernelRegistry() + spec = self._spec() + reg.register(spec) + self.assertEqual(reg.get("g.n"), [spec]) + self.assertTrue(reg.has("g.n")) + self.assertEqual(reg.ops(), ["g.n"]) + + def test_get_unknown_op_returns_empty(self): + reg = KernelRegistry() + self.assertEqual(reg.get("no.such"), []) + self.assertFalse(reg.has("no.such")) + + def test_reregister_same_backend_replaces(self): + reg = KernelRegistry() + reg.register(self._spec(target="math:sqrt")) + reg.register(self._spec(target="math:floor")) + specs = reg.get("g.n") + self.assertEqual(len(specs), 1) + self.assertEqual(specs[0].target, "math:floor") + + def test_get_backend_missing_raises(self): + reg = KernelRegistry() + reg.register(self._spec(backend=KernelBackend.TORCH)) + with self.assertRaises(KeyError): + reg.get_backend("g.n", KernelBackend.TRITON) + with self.assertRaises(KeyError): + reg.get_backend("no.such", KernelBackend.TORCH) + + +class TestKernelSpecUnit(unittest.TestCase): + def test_load_simple_target(self): + import math + + spec = KernelSpec(op="g.n", backend=KernelBackend.TORCH, target="math:sqrt") + self.assertIs(spec.load(), math.sqrt) + + def test_load_dotted_target(self): + spec = KernelSpec( + op="g.n", + backend=KernelBackend.TORCH, + target="sglang.kernels.ops.layernorm:_RMSNORM.forward_native", + ) + self.assertTrue(callable(spec.load())) + + def test_load_bad_target_raises(self): + spec = KernelSpec(op="g.n", backend=KernelBackend.TORCH, target="no-colon") + with self.assertRaises(ValueError): + spec.load() + + +def _ref_rmsnorm(x, w, eps): + xf = x.to(torch.float32) + var = xf.pow(2).mean(dim=-1, keepdim=True) + return (xf * torch.rsqrt(var + eps) * w).to(x.dtype) + + +class TestNativeReferenceImplementations(unittest.TestCase): + """The forward_native math of the reworked ops, on CPU tensors.""" + + def setUp(self): + torch.manual_seed(0) + + def test_rmsnorm_native(self): + from sglang.kernels.ops.layernorm import _RMSNORM + + x = torch.randn(8, 128) + w = torch.randn(128) + out = _RMSNORM.forward_native(x, w, 1e-6) + self.assertTrue(torch.allclose(out, _ref_rmsnorm(x, w, 1e-6))) + # out= writes in place and returns out + buf = torch.empty_like(x) + self.assertIs(_RMSNORM.forward_native(x, w, 1e-6, out=buf), buf) + self.assertTrue(torch.allclose(buf, out)) + + def test_fused_add_rmsnorm_native(self): + from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM + + x = torch.randn(8, 128) + residual = torch.randn(8, 128) + w = torch.randn(128) + x2, r2 = x.clone(), residual.clone() + self.assertIsNone(_FUSED_ADD_RMSNORM.forward_native(x, residual, w, 1e-6)) + acc = x2.to(torch.float32) + r2.to(torch.float32) + self.assertTrue(torch.allclose(residual, acc)) + ref = acc * torch.rsqrt(acc.pow(2).mean(-1, keepdim=True) + 1e-6) * w + self.assertTrue(torch.allclose(x, ref)) + + def test_gemma_rmsnorm_native(self): + from sglang.kernels.ops.layernorm import _GEMMA_RMSNORM + + x = torch.randn(8, 128) + w = torch.randn(128) + out = _GEMMA_RMSNORM.forward_native(x, w, 1e-6) + xf = x.to(torch.float32) + ref = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6) * (1.0 + w) + self.assertTrue(torch.allclose(out, ref)) + + def test_gated_activations_native(self): + import torch.nn.functional as F + + from sglang.kernels.ops.activation import ( + _GELU_AND_MUL, + _GELU_TANH_AND_MUL, + _SILU_AND_MUL, + ) + + x = torch.randn(8, 256) + gate, up = x[..., :128], x[..., 128:] + cases = [ + (_SILU_AND_MUL, F.silu(gate) * up), + (_GELU_AND_MUL, F.gelu(gate, approximate="none") * up), + (_GELU_TANH_AND_MUL, F.gelu(gate, approximate="tanh") * up), + ] + for op, ref in cases: + self.assertTrue(torch.allclose(op.forward_native(x), ref), op.op) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/test_fused_op_gpu_parity.py b/test/registered/kernels/test_fused_op_gpu_parity.py new file mode 100644 index 000000000..1178748f7 --- /dev/null +++ b/test/registered/kernels/test_fused_op_gpu_parity.py @@ -0,0 +1,125 @@ +"""Generic every-backend-vs-native parity harness for BaseFusedOp operators. + +Part of RFC #29630, Phase 2. For each reworked fused op, enumerate its +available backends, run each one that is eligible on this platform, and +assert the output matches the pure-torch ``forward_native`` reference within +dtype tolerance. New backends added to an op are picked up automatically — +no per-kernel test boilerplate. +""" + +import unittest + +import torch + +from sglang.kernels.spec import KernelBackend +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + +_DEVICE = "cuda" +# torch_compile is native under the hood; exclude it from the sweep to keep +# CI time down (compilation dominates) — it is exercised in the CPU lane. +_SKIP_BACKENDS = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE} +_TOLERANCE = { + torch.float16: dict(atol=1e-2, rtol=1e-2), + torch.bfloat16: dict(atol=2e-2, rtol=2e-2), +} + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +class TestFusedOpGpuParity(CustomTestCase): + def setUp(self): + torch.manual_seed(0) + + def _eligible_backends(self, op): + return [ + b + for b in op.available_backends() + if b not in _SKIP_BACKENDS and op.backend_eligible(b) + ] + + def _assert_close(self, got, ref, dtype, msg): + torch.testing.assert_close(got, ref, **_TOLERANCE[dtype], msg=msg) + + def test_rmsnorm_backends_match_native(self): + from sglang.kernels.ops.layernorm import _RMSNORM + + for dtype in (torch.float16, torch.bfloat16): + for shape in ((1, 4096), (128, 4096), (7, 2048)): + x = torch.randn(shape, dtype=dtype, device=_DEVICE) + w = torch.randn(shape[-1], dtype=dtype, device=_DEVICE) + ref = _RMSNORM.forward_native(x, w, 1e-6) + for backend in self._eligible_backends(_RMSNORM): + got = _RMSNORM.forward(x, w, 1e-6, backend=backend) + self._assert_close( + got, ref, dtype, f"rmsnorm {backend.value} {dtype} {shape}" + ) + + def test_fused_add_rmsnorm_backends_match_native(self): + from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM + + for dtype in (torch.float16, torch.bfloat16): + for shape in ((1, 4096), (128, 4096)): + x0 = torch.randn(shape, dtype=dtype, device=_DEVICE) + r0 = torch.randn(shape, dtype=dtype, device=_DEVICE) + w = torch.randn(shape[-1], dtype=dtype, device=_DEVICE) + x_ref, r_ref = x0.clone(), r0.clone() + _FUSED_ADD_RMSNORM.forward_native(x_ref, r_ref, w, 1e-6) + for backend in self._eligible_backends(_FUSED_ADD_RMSNORM): + x, r = x0.clone(), r0.clone() + _FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=backend) + label = f"fused_add_rmsnorm {backend.value} {dtype} {shape}" + self._assert_close(x, x_ref, dtype, label + " (normed)") + self._assert_close(r, r_ref, dtype, label + " (residual)") + + def test_gemma_rmsnorm_backends_match_native(self): + from sglang.kernels.ops.layernorm import _GEMMA_RMSNORM + + for dtype in (torch.float16, torch.bfloat16): + x = torch.randn(64, 2048, dtype=dtype, device=_DEVICE) + w = torch.randn(2048, dtype=dtype, device=_DEVICE) + ref = _GEMMA_RMSNORM.forward_native(x, w, 1e-6) + for backend in self._eligible_backends(_GEMMA_RMSNORM): + got = _GEMMA_RMSNORM.forward(x, w, 1e-6, backend=backend) + self._assert_close( + got, ref, dtype, f"gemma_rmsnorm {backend.value} {dtype}" + ) + + def test_gemma_fused_add_rmsnorm_backends_match_native(self): + from sglang.kernels.ops.layernorm import _GEMMA_FUSED_ADD_RMSNORM + + for dtype in (torch.float16, torch.bfloat16): + x0 = torch.randn(64, 2048, dtype=dtype, device=_DEVICE) + r0 = torch.randn(64, 2048, dtype=dtype, device=_DEVICE) + w = torch.randn(2048, dtype=dtype, device=_DEVICE) + x_ref, r_ref = x0.clone(), r0.clone() + _GEMMA_FUSED_ADD_RMSNORM.forward_native(x_ref, r_ref, w, 1e-6) + for backend in self._eligible_backends(_GEMMA_FUSED_ADD_RMSNORM): + x, r = x0.clone(), r0.clone() + _GEMMA_FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=backend) + label = f"gemma_fused_add_rmsnorm {backend.value} {dtype}" + self._assert_close(x, x_ref, dtype, label + " (normed)") + self._assert_close(r, r_ref, dtype, label + " (residual)") + + def test_gated_activation_backends_match_native(self): + from sglang.kernels.ops.activation import ( + _GELU_AND_MUL, + _GELU_TANH_AND_MUL, + _SILU_AND_MUL, + ) + + for op in (_SILU_AND_MUL, _GELU_AND_MUL, _GELU_TANH_AND_MUL): + for dtype in (torch.float16, torch.bfloat16): + for shape in ((1, 8192), (128, 8192)): + x = torch.randn(shape, dtype=dtype, device=_DEVICE) + ref = op.forward_native(x) + for backend in self._eligible_backends(op): + got = op.forward(x, backend=backend) + self._assert_close( + got, ref, dtype, f"{op.op} {backend.value} {dtype} {shape}" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/test_gather_spec_extras.py b/test/registered/kernels/test_gather_spec_extras.py index a2ad11fb3..f22e5498d 100644 --- a/test/registered/kernels/test_gather_spec_extras.py +++ b/test/registered/kernels/test_gather_spec_extras.py @@ -7,7 +7,7 @@ import unittest import torch -from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras +from sglang.kernels.ops.speculative.gather_spec_extras import gather_spec_extras from sglang.test.test_utils import CustomTestCase _OUTPUT_NAMES = ("topk_p", "topk_index", "bonus_tokens", "hidden_states") diff --git a/test/registered/kernels/test_kernels_namespace.py b/test/registered/kernels/test_kernels_namespace.py new file mode 100644 index 000000000..954ca94c9 --- /dev/null +++ b/test/registered/kernels/test_kernels_namespace.py @@ -0,0 +1,256 @@ +"""GPU-free import/registry tests for the ``sglang.kernels`` namespace. + +Part of RFC #29630, Phase 2. These tests exercise the public namespace, the +kernel registry, and the heuristic selector without touching a GPU or importing +any kernel backend (``sgl_kernel`` / ``sglang.jit_kernel``). They run in the CPU +CI lane. +""" + +import subprocess +import sys +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +# A must-contain subset of registered operators and their backends. The +# registry holds many more entries (every migrated Triton kernel), so this is +# checked as a subset, not an exact match. +EXPECTED_OPS = { + # BaseFusedOp-backed ops: native + torch_compile always available, + # plus the overridden CUDA backends. + "activation.silu_and_mul": {"cuda_aot", "cuda_jit", "torch", "torch_compile"}, + "activation.gelu_and_mul": {"cuda_aot", "cuda_jit", "torch", "torch_compile"}, + "activation.gelu_tanh_and_mul": { + "cuda_aot", + "cuda_jit", + "torch", + "torch_compile", + }, + "layernorm.rmsnorm": {"cuda_aot", "cuda_jit", "torch", "torch_compile"}, + "layernorm.fused_add_rmsnorm": { + "cuda_aot", + "cuda_jit", + "torch", + "torch_compile", + }, + "layernorm.gemma_rmsnorm": {"cuda_aot", "torch", "torch_compile"}, + "layernorm.gemma_fused_add_rmsnorm": {"cuda_aot", "torch", "torch_compile"}, + # curated dual/single-backend wrapper ops + "gemm.fp8_scaled_mm": {"cuda_aot"}, + "gemm.dsv3_fused_a_gemm": {"cuda_aot", "cuda_jit"}, + "gemm.dsv3_router_gemm": {"cuda_jit"}, + "kvcache.reshape_and_cache_flash": {"triton"}, + "moe.moe_align_block_size": {"cuda_aot", "cuda_jit"}, + "moe.topk_softmax": {"cuda_aot"}, + "quantization.sgl_per_token_quant_fp8": {"cuda_aot"}, + "quantization.sgl_per_token_group_quant_8bit": {"cuda_aot", "cuda_jit"}, + "quantization.sgl_per_token_group_quant_fp8": {"cuda_aot"}, + "quantization.sgl_per_token_group_quant_int8": {"cuda_aot"}, + # deferred-group wrappers, now populated + "sampling.top_k_renorm_probs": {"cuda_aot"}, + "sampling.top_p_renorm_probs": {"cuda_aot"}, + "spatial.get_sm_available": {"cuda_aot"}, + "spatial.create_greenctx_stream_by_value": {"cuda_aot"}, + "mamba.causal_conv1d_fwd": {"cuda_aot"}, + "mamba.causal_conv1d_update": {"cuda_aot"}, + "diffusion.apply_group_norm_silu": {"cuda_jit"}, + "diffusion.residual_gate_add": {"cuda_jit"}, + "diffusion.fused_inplace_qknorm_rope": {"cuda_jit"}, + # representative migrated Triton kernels (inventory) + "grammar.apply_token_bitmask_inplace_triton": {"triton"}, + "memory.alloc_extend_kernel": {"triton"}, + "attention.decode_attention_fwd": {"triton"}, + "kvcache.create_flashinfer_kv_indices_triton": {"triton"}, + "speculative.gather_spec_extras": {"triton"}, +} + +# Public wrapper callables that each populated group must expose. +EXPECTED_WRAPPERS = { + "sglang.kernels.ops.layernorm": [ + "rmsnorm", + "fused_add_rmsnorm", + "gemma_rmsnorm", + "gemma_fused_add_rmsnorm", + ], + "sglang.kernels.ops.activation": [ + "silu_and_mul", + "gelu_and_mul", + "gelu_tanh_and_mul", + ], + "sglang.kernels.ops.gemm": [ + "fp8_scaled_mm", + "dsv3_fused_a_gemm", + "dsv3_router_gemm", + ], + "sglang.kernels.ops.quantization": [ + "sgl_per_token_quant_fp8", + "sgl_per_token_group_quant_8bit", + "sgl_per_token_group_quant_fp8", + "sgl_per_token_group_quant_int8", + ], + "sglang.kernels.ops.moe": ["moe_align_block_size", "topk_softmax"], + "sglang.kernels.ops.kvcache": ["reshape_and_cache_flash"], + "sglang.kernels.ops.sampling": ["top_k_renorm_probs", "top_p_renorm_probs"], + "sglang.kernels.ops.spatial": [ + "get_sm_available", + "create_greenctx_stream_by_value", + ], + "sglang.kernels.ops.mamba": ["causal_conv1d_fwd", "causal_conv1d_update"], + "sglang.kernels.ops.diffusion": [ + "apply_group_norm_silu", + "residual_gate_add", + "fused_inplace_qknorm_rope", + ], +} + +# All operator groups from the RFC's proposed shape must import as packages. +ALL_GROUPS = [ + "activation", + "attention", + "communication", + "diffusion", + "gemm", + "grammar", + "kvcache", + "layernorm", + "mamba", + "memory", + "moe", + "quantization", + "sampling", + "spatial", + "speculative", +] + + +class TestKernelsNamespace(unittest.TestCase): + def setUp(self): + import importlib + + import sglang.kernels + import sglang.kernels.ops # populate the registry + + self.K = sglang.kernels + self.importlib = importlib + + def test_top_level_exports(self): + for name in ( + "KernelSpec", + "KernelBackend", + "FormatSignature", + "CapabilityRequirement", + "PlatformInfo", + "registry", + "get_kernel", + "select_kernel", + ): + self.assertTrue(hasattr(self.K, name), f"missing export: {name}") + + def test_all_groups_importable(self): + for group in ALL_GROUPS: + mod = self.importlib.import_module(f"sglang.kernels.ops.{group}") + self.assertTrue(hasattr(mod, "__all__")) + + def test_registry_contents(self): + registry = self.K.registry + ops = set(registry.ops()) + # EXPECTED_OPS is a must-contain subset (many more migrated kernels + # are also registered). + missing = set(EXPECTED_OPS) - ops + self.assertFalse(missing, f"missing registered ops: {sorted(missing)}") + for op, backends in EXPECTED_OPS.items(): + got = {s.backend.value for s in registry.get(op)} + self.assertEqual(got, backends, f"backend mismatch for {op}") + self.assertGreaterEqual(len(ops), 80, "registry unexpectedly small") + + def test_specs_are_well_formed(self): + for spec in self.K.registry.all_specs(): + self.assertIn(".", spec.op) + self.assertEqual(spec.op, f"{spec.group}.{spec.name}") + # target must be an importable "module:attr" path + module_path, sep, attr = spec.target.partition(":") + self.assertEqual(sep, ":", f"bad target for {spec.op}: {spec.target}") + self.assertTrue(module_path and attr, spec.target) + + def test_wrappers_exposed_and_callable(self): + for module_name, names in EXPECTED_WRAPPERS.items(): + mod = self.importlib.import_module(module_name) + for name in names: + self.assertTrue(callable(getattr(mod, name)), f"{module_name}.{name}") + + def test_single_backend_op_resolves_without_backend(self): + # An op with exactly one registered backend has a fixed call path. + for op, backends in EXPECTED_OPS.items(): + if len(backends) == 1: + spec = self.K.select_kernel(op) + self.assertEqual(spec.backend.value, next(iter(backends)), op) + + def test_multi_backend_op_requires_explicit_backend(self): + # No hidden ranking: a multi-backend op must be resolved explicitly. + multi = [op for op, b in EXPECTED_OPS.items() if len(b) > 1] + self.assertTrue(multi) # sanity: we do have multi-backend ops + for op in multi: + with self.assertRaises(ValueError): + self.K.select_kernel(op) + + def test_selector_explicit_backend(self): + spec = self.K.select_kernel( + "layernorm.rmsnorm", backend=self.K.KernelBackend.CUDA_JIT + ) + self.assertEqual( + spec.target, "sglang.kernels.ops.layernorm:_RMSNORM.forward_cuda_jit" + ) + + def test_selector_unknown_op_raises(self): + with self.assertRaises(KeyError): + self.K.select_kernel("does_not.exist") + with self.assertRaises(KeyError): + self.K.select_kernel( + "gemm.fp8_scaled_mm", backend=self.K.KernelBackend.TRITON + ) + + def test_capability_requirement_logic(self): + cap = self.K.CapabilityRequirement + plat = self.K.PlatformInfo + cpu = plat(device_type="cpu") + sm90 = plat(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0) + sm100 = plat(device_type="cuda", cuda_arch_major=10, cuda_arch_minor=0) + + self.assertFalse(cap(requires_cuda=True).is_satisfied_by(cpu)) + self.assertTrue(cap(requires_cuda=True).is_satisfied_by(sm90)) + self.assertFalse( + cap(requires_cuda=True, min_cuda_arch=(10, 0)).is_satisfied_by(sm90) + ) + self.assertTrue( + cap(requires_cuda=True, min_cuda_arch=(10, 0)).is_satisfied_by(sm100) + ) + self.assertFalse( + cap(requires_cuda=True, max_cuda_arch=(9, 0)).is_satisfied_by(sm100) + ) + + def test_platform_detect_does_not_raise(self): + plat = self.K.PlatformInfo.detect() + self.assertIn(plat.device_type, ("cpu", "cuda", "hip")) + + def test_import_does_not_load_kernel_backends(self): + # Importing the namespace must stay metadata-only: no sgl_kernel or + # sglang.jit_kernel import, and no JIT compilation, on a CPU box. + code = ( + "import sys; import sglang.kernels.ops; " + "backend = ('sgl_kernel' in sys.modules) or " + "any(m.startswith('sglang.jit_kernel') for m in sys.modules); " + "print('BACKEND_IMPORTED' if backend else 'CLEAN')" + ) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("CLEAN", result.stdout, result.stdout + result.stderr) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/lora/test_chunked_sgmv_backend.py b/test/registered/lora/test_chunked_sgmv_backend.py index e075043a7..7658bf1f9 100644 --- a/test/registered/lora/test_chunked_sgmv_backend.py +++ b/test/registered/lora/test_chunked_sgmv_backend.py @@ -5,19 +5,25 @@ from typing import List, Optional, Tuple import torch -from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor -from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend -from sglang.srt.lora.triton_ops import ( +from sglang.kernels.ops.gemm.chunked_embedding_lora_a import ( chunked_embedding_lora_a_forward, +) +from sglang.kernels.ops.gemm.chunked_sgmv_expand import ( + _chunked_lora_expand_kernel, chunked_sgmv_lora_expand_forward, +) +from sglang.kernels.ops.gemm.chunked_sgmv_shrink import ( + _chunked_lora_shrink_kernel, chunked_sgmv_lora_shrink_forward, +) +from sglang.kernels.ops.gemm.kv_b_lora_absorbed import ( step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, step_b_v_fwd, ) -from sglang.srt.lora.triton_ops.chunked_sgmv_expand import _chunked_lora_expand_kernel -from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel +from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor +from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend from sglang.srt.lora.utils import LoRABatchInfo, get_lm_head_pruned_lens from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/lora/test_fused_moe_lora_kernel.py b/test/registered/lora/test_fused_moe_lora_kernel.py index 9b152464c..1cae10d8d 100644 --- a/test/registered/lora/test_fused_moe_lora_kernel.py +++ b/test/registered/lora/test_fused_moe_lora_kernel.py @@ -9,7 +9,7 @@ import torch # IMPORT PREBUILT KERNEL # ============================================================================== from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size -from sglang.srt.lora.triton_ops import fused_moe_lora +from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora from sglang.srt.utils import set_random_seed from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/lora/test_virtual_experts_kernels.py b/test/registered/lora/test_virtual_experts_kernels.py index 91423abc8..29c5b8a27 100644 --- a/test/registered/lora/test_virtual_experts_kernels.py +++ b/test/registered/lora/test_virtual_experts_kernels.py @@ -32,7 +32,7 @@ from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small") -from sglang.srt.lora.triton_ops.virtual_experts import ( +from sglang.kernels.ops.moe.virtual_experts import ( _align_block_size_jit, _align_block_size_torch, _fused_virtual_topk_ids, diff --git a/test/registered/unit/constrained/test_token_filter_ops.py b/test/registered/unit/constrained/test_token_filter_ops.py index 363f6a818..47a9818cd 100644 --- a/test/registered/unit/constrained/test_token_filter_ops.py +++ b/test/registered/unit/constrained/test_token_filter_ops.py @@ -19,7 +19,7 @@ register_cpu_ci(2.0, "base-a-test-cpu") # Conditionally import Triton path _has_cuda = torch.cuda.is_available() if _has_cuda: - from sglang.srt.constrained.triton_ops.token_filter_ops import ( + from sglang.kernels.ops.grammar.token_filter_ops import ( set_token_filter_triton, ) diff --git a/test/registered/unit/mem_cache/test_store_cache_4d.py b/test/registered/unit/mem_cache/test_store_cache_4d.py index ba874fe18..d899f1b0c 100644 --- a/test/registered/unit/mem_cache/test_store_cache_4d.py +++ b/test/registered/unit/mem_cache/test_store_cache_4d.py @@ -111,7 +111,7 @@ class TestStoreCache4D(unittest.TestCase): dtype: torch.dtype = torch.bfloat16, loc_dtype: torch.dtype = torch.int64, ): - from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d + from sglang.kernels.ops.kvcache.cache_move import store_cache_4d # Two independent target buffers — one for the kernel, one for the # legacy reference path. @@ -219,7 +219,7 @@ class TestStoreCache4D(unittest.TestCase): def test_store_cache_4d_empty_loc(self): """N=0 must be a no-op: no kernel launch, no exception, no buffer mutation.""" - from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d + from sglang.kernels.ops.kvcache.cache_move import store_cache_4d k_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda") v_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda") @@ -259,7 +259,7 @@ class TestStoreCache4DAssertions(unittest.TestCase): """Wrapper requires `stride[-1] == 1` and `stride[-2] == head_dim` (the trailing two dims must be contiguous). A permutation that breaks this should trigger AssertionError.""" - from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d + from sglang.kernels.ops.kvcache.cache_move import store_cache_4d # Build a 4-D view, then permute the last two dims → trailing # contiguity violated. @@ -278,7 +278,7 @@ class TestStoreCache4DAssertions(unittest.TestCase): def test_rejects_dtype_mismatch(self): """All four tensors must share a dtype; the caller is responsible for any cast before the call.""" - from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d + from sglang.kernels.ops.kvcache.cache_move import store_cache_4d k_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda") v_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda") diff --git a/test/registered/unit/mem_cache/test_triton_kernel_layout.py b/test/registered/unit/mem_cache/test_triton_kernel_layout.py index f881c6ca8..9d4b56816 100644 --- a/test/registered/unit/mem_cache/test_triton_kernel_layout.py +++ b/test/registered/unit/mem_cache/test_triton_kernel_layout.py @@ -57,7 +57,7 @@ class TestTritonKernelLayoutParity(unittest.TestCase): return q, logical_kv_k, logical_kv_v, kv_indptr, kv_indices, seq_len def _run_decode(self, q, k_buf, v_buf, kv_indptr, kv_indices, page_size): - from sglang.srt.layers.attention.triton_ops.decode_attention import ( + from sglang.kernels.ops.attention.decode_attention import ( decode_attention_fwd, ) @@ -115,7 +115,7 @@ class TestTritonKernelLayoutParity(unittest.TestCase): def test_extend_3d_vs_4d_ps1_byte_identical(self): """Same parity check for extend kernel.""" - from sglang.srt.layers.attention.triton_ops.extend_attention import ( + from sglang.kernels.ops.attention.extend_attention import ( extend_attention_fwd, )