chore: import upstream snapshot with attribution

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/
+    <group>/     # 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.<group>` 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
+  (`"<group>.<name>"`), 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_<backend>` 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.<group>` 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.

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.<group>``, 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_<backend>`` 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",
+]

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_<backend>`` 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_<backend>`` 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_<backend> 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_<backend>`` 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, ``"<group>.<name>"`` (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.
+    ``"<module>:<attr>.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

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)

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

python/sglang/kernels/ops/activation/softcap.py

@@ -0,0 +1,120 @@
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra import libdevice
+
+softcap_out_autotune = triton.autotune(
+    configs=[
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 32768}, num_warps=32),
+    ],
+    key=["n_ele"],
+)
+
+
+@triton.jit
+def softcap_out_kernel(
+    output_ptr,
+    input_ptr,
+    n_ele,
+    softcap_const: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_ele
+    x = tl.load(input_ptr + offsets, mask=mask)
+    fx = x.to(tl.float32)
+    fxs = fx / softcap_const
+    exped = tl.exp(2 * fxs)
+    top = exped - 1
+    bottom = exped + 1
+    output = top / bottom * softcap_const
+    tl.store(output_ptr + offsets, output, mask=mask)
+
+
+softcap_out_kernel_autotuned = softcap_out_autotune(softcap_out_kernel)
+
+
+def softcap_out(x, softcap_const, autotune=False):
+    output = torch.empty_like(x, dtype=torch.float32)
+    n_elements = output.numel()
+    if autotune:
+
+        def grid(meta):
+            return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+        softcap_out_kernel_autotuned[grid](output, x, n_elements, softcap_const)
+    else:
+        softcap_out_kernel[(triton.cdiv(n_elements, 128),)](
+            output, x, n_elements, softcap_const, BLOCK_SIZE=128, num_warps=8
+        )
+    return output
+
+
+@triton.jit
+def softcap_inplace_logits_kernel(
+    full_logits_ptr,
+    softcapping_value,
+    ncols,
+    row_stride,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(1).to(tl.int64)
+    pid = tl.program_id(0).to(tl.int64)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < ncols
+
+    # Load values
+    row_ptr = full_logits_ptr + row * row_stride
+    x = tl.load(row_ptr + offsets, mask=mask)
+
+    # Perform operations in-place
+    x = x / softcapping_value
+    x = libdevice.tanh(x)
+    x = x * softcapping_value
+
+    # Store result
+    tl.store(row_ptr + offsets, x, mask=mask)
+
+
+def softcap_inplace_logits(full_logits, final_logit_softcapping):
+    if full_logits.is_contiguous():
+        nrows, ncols = 1, full_logits.numel()
+        row_stride = ncols
+    else:
+        assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor"
+        assert (
+            full_logits.stride(1) == 1
+        ), "non-contiguous softcap requires contiguous columns"
+        nrows, ncols = full_logits.shape
+        row_stride = full_logits.stride(0)
+
+    BLOCK_SIZE = 1024
+    grid = ((ncols + BLOCK_SIZE - 1) // BLOCK_SIZE, nrows)
+
+    softcap_inplace_logits_kernel[grid](
+        full_logits_ptr=full_logits,
+        softcapping_value=final_logit_softcapping,
+        ncols=ncols,
+        row_stride=row_stride,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return full_logits

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.<module>``); 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__ = []

python/sglang/kernels/ops/attention/decode_attention.py

@@ -0,0 +1,969 @@
+# 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.
+# ==============================================================================
+"""
+Memory-efficient attention for decoding.
+It supports page size = 1.
+"""
+
+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
+
+import logging
+
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+logger = logging.getLogger(__name__)
+
+
+_MIN_BLOCK_KV = 32
+
+
+def _extract_kv_strides(buf, page_size: int):
+    """Extract (slot_stride, head_stride, page_stride, tok_stride) for a
+    KV buffer that may be:
+      - 3-D ``[max_slots, head_num, head_dim]`` (legacy / non-shared) — the
+        contiguous layout most callers use. page/tok strides are synthesized
+        so the kernel's PAGE_SIZE>1 math collapses to ``kv_loc * stride(0)``.
+      - 4-D ``[num_pages, page_size, head_num, head_dim]`` (shared
+        pool). page/tok strides come from stride(0)/stride(1) directly;
+        legacy ``stride_bs`` is set to 0 (unused at PAGE_SIZE>1).
+
+    Returns a 4-tuple of ints suitable for passing as ``stride_buf_*bs``,
+    ``stride_buf_*h``, ``stride_buf_*page``, ``stride_buf_*tok``.
+    """
+    if buf.ndim == 4:
+        # 4-D view ``[num_pages, page_size, head_num, head_dim]``.
+        #   stride(0) = per-PAGE stride (page_bytes/itemsize)
+        #   stride(1) = within-page per-TOKEN stride (k_row/v_row bytes/itemsize)
+        # The PAGE_SIZE>1 kernel branch uses page_stride/tok_stride and does
+        # NOT read slot_stride. slot_stride is consumed ONLY by the
+        # PAGE_SIZE==1 branch (``offs = kv_loc * stride_buf_*bs``), where one
+        # page holds exactly one slot, so the per-slot stride is the per-page
+        # stride — NOT the within-page token stride. Concretely the per-slot
+        # stride is ``page_stride // page_size`` (= entry_bytes/itemsize),
+        # which at ps=1 equals page_stride. Using ``tok_stride`` here (one
+        # layer's k_row) would make the ps=1 read address ``kv_loc * k_row``
+        # instead of ``kv_loc * entry_bytes`` and read the wrong slot.
+        page_stride = buf.stride(0)
+        tok_stride = buf.stride(1)
+        head_stride = buf.stride(2)
+        slot_stride = (
+            page_stride // page_size
+        )  # per-slot stride; == page_stride at ps=1
+        assert buf.shape[1] == page_size, (
+            f"4-D KV buffer's dim-1 must equal page_size; got "
+            f"shape[1]={buf.shape[1]}, page_size={page_size}"
+        )
+    elif buf.ndim == 3:
+        # Legacy 3-D ``[N, head, dim]``. Synthesize page/tok strides such
+        # that ``(kv_loc // ps) * page_stride + (kv_loc % ps) * tok_stride
+        # == kv_loc * slot_stride`` for the page-aware branch — this lets
+        # the same kernel handle non-shared paged-allocator buffers without
+        # any caller adjustment.
+        slot_stride = buf.stride(0)
+        head_stride = buf.stride(1)
+        page_stride = slot_stride * page_size
+        tok_stride = slot_stride
+    else:  # pragma: no cover
+        raise ValueError(f"unexpected KV buffer ndim={buf.ndim}, shape={buf.shape}")
+    return slot_stride, head_stride, page_stride, tok_stride
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.jit
+def _fwd_kernel_stage1(
+    Q,
+    K_Buffer,
+    V_Buffer,
+    sm_scale_withk,
+    kv_indptr,
+    kv_indices,
+    Att_Out,
+    Att_Lse,
+    num_kv_splits,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    # Page-aware strides (used when PAGE_SIZE > 1). For
+    # PAGE_SIZE == 1 the address math degenerates and these are unused
+    # (Triton specializes the dead branch away at compile time).
+    stride_buf_kpage,
+    stride_buf_ktok,
+    stride_buf_vpage,
+    stride_buf_vtok,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    kv_group_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    logit_cap: tl.constexpr,
+    Lk: tl.constexpr,
+    Lv: tl.constexpr,
+    xai_temperature_len: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+):
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
+    cur_head = tl.program_id(1)
+    split_kv_id = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lk
+    mask_dv = offs_dv < Lv
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    if xai_temperature_len > 0:
+        offs_qidx = cur_batch_seq_len - 1
+        xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len))
+        _qtemp = tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale
+        xai_temperature_reg = tl.where(offs_qidx > xai_temperature_len, _qtemp, 1.0)
+
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
+
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = -float("inf")
+    e_sum = 0.0
+    acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        q = tl.load(Q + off_q, mask=mask_d, other=0.0)
+        for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+            # Page-aware KV address math. At PAGE_SIZE==1 (legacy
+            # / non-shared / shared-at-ps=1), Triton specializes the
+            # else-branch away and the SASS is byte-identical to today.
+            if PAGE_SIZE == 1:
+                offs_buf_k = (
+                    kv_loc[:, None] * stride_buf_kbs
+                    + cur_kv_head * stride_buf_kh
+                    + offs_d[None, :]
+                )
+            else:
+                page_id = kv_loc // PAGE_SIZE
+                tok_in_p = kv_loc % PAGE_SIZE
+                offs_buf_k = (
+                    page_id[:, None] * stride_buf_kpage
+                    + tok_in_p[:, None] * stride_buf_ktok
+                    + cur_kv_head * stride_buf_kh
+                    + offs_d[None, :]
+                )
+            k = tl.load(
+                K_Buffer + offs_buf_k,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]),
+                other=0.0,
+            )
+            qk = tl.sum(q[None, :] * k, 1)
+            qk *= sm_scale_withk
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            if xai_temperature_len > 0:
+                qk *= xai_temperature_reg
+
+            qk = tl.where(offs_n < split_kv_end, qk, float("-inf"))
+
+            if PAGE_SIZE == 1:
+                offs_buf_v = (
+                    kv_loc[:, None] * stride_buf_vbs
+                    + cur_kv_head * stride_buf_vh
+                    + offs_dv[None, :]
+                )
+            else:
+                offs_buf_v = (
+                    page_id[:, None] * stride_buf_vpage
+                    + tok_in_p[:, None] * stride_buf_vtok
+                    + cur_kv_head * stride_buf_vh
+                    + offs_dv[None, :]
+                )
+            v = tl.load(
+                V_Buffer + offs_buf_v,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
+                other=0.0,
+            )
+
+            n_e_max = tl.maximum(tl.max(qk, 0), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max)
+            acc *= re_scale
+            acc += tl.sum(p[:, None] * v, 0)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 0)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_dv
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum,
+            mask=(mask_dv),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+        ) // Lv
+
+        tl.store(
+            Att_Lse + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+        )
+
+
+def _decode_att_m_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    att_lse,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    logit_cap,
+    xai_temperature_len=-1,
+    page_size: int = 1,
+):
+    BLOCK = 64
+    # [TODO] work around SGPR limit on MI3xx
+    if _is_hip:
+        BLOCK = 8
+    MAX_KV_SPLITS = max_kv_splits
+    Lk = k_buffer.shape[-1]
+    Lv = v_buffer.shape[-1]
+
+    # head_num lives in the dim immediately before the head_dim. For 3-D
+    # ``[N, head_num, head_dim]`` that's dim 1; for 4-D
+    # ``[num_pages, page_size, head_num, head_dim]`` that's dim 2.
+    kv_head_num = k_buffer.shape[-2]
+
+    batch, head_num = q.shape[0], q.shape[1]
+
+    grid = (batch, head_num, MAX_KV_SPLITS)
+    kv_group_num = q.shape[1] // kv_head_num
+
+    if kv_group_num == 1:
+        num_warps = 4
+    else:
+        num_warps = 2
+        if _is_hip:
+            num_warps = 1
+
+    BLOCK_DMODEL = triton.next_power_of_2(Lk)
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    k_slot_stride, k_head_stride, k_page_stride, k_tok_stride = _extract_kv_strides(
+        k_buffer, page_size
+    )
+    v_slot_stride, v_head_stride, v_page_stride, v_tok_stride = _extract_kv_strides(
+        v_buffer, page_size
+    )
+
+    _fwd_kernel_stage1[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        sm_scale_withk,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        att_lse,
+        num_kv_splits,
+        q.stride(0),
+        q.stride(1),
+        k_slot_stride,
+        k_head_stride,
+        v_slot_stride,
+        v_head_stride,
+        k_page_stride,
+        k_tok_stride,
+        v_page_stride,
+        v_tok_stride,
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        kv_group_num=kv_group_num,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_N=BLOCK,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        logit_cap=logit_cap,
+        xai_temperature_len=xai_temperature_len,
+        num_warps=num_warps,
+        num_stages=2,
+        Lk=Lk,
+        Lv=Lv,
+        PAGE_SIZE=page_size,
+    )
+
+
+@triton.jit
+def _fwd_grouped_kernel_stage1(
+    Q,
+    K_Buffer,
+    V_Buffer,
+    sm_scale_withk,
+    kv_indptr,
+    kv_indices,
+    Att_Out,
+    Att_Lse,
+    num_kv_splits,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    # Page-aware strides (used when PAGE_SIZE > 1).
+    stride_buf_kpage,
+    stride_buf_ktok,
+    stride_buf_vpage,
+    stride_buf_vtok,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    kv_group_num: tl.constexpr,
+    q_head_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    logit_cap: tl.constexpr,
+    xai_temperature_len: tl.constexpr,
+    Lk: tl.constexpr,
+    Lv: tl.constexpr,
+    HAS_MLA: tl.constexpr = False,
+    USE_PDL: tl.constexpr = False,
+    PAGE_SIZE: tl.constexpr = 1,
+):
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
+    cur_head_id = tl.program_id(1)
+    cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
+    split_kv_id = tl.program_id(2)
+
+    if BLOCK_H < kv_group_num:
+        VALID_BLOCK_H: tl.constexpr = BLOCK_H
+    else:
+        VALID_BLOCK_H: tl.constexpr = kv_group_num
+    cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
+    mask_h = mask_h & (cur_head < q_head_num)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lk
+    mask_dv = offs_dv < Lv
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    if xai_temperature_len > 0:
+        offs_qidx = cur_batch_seq_len - 1
+        xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len))
+        _qtemp = tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale
+        xai_temperature_reg = tl.where(offs_qidx > xai_temperature_len, _qtemp, 1.0)
+
+    offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :]
+
+    if BLOCK_DPE > 0:
+        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
+        mask_dpe = offs_dpe < Lk
+        off_qpe = (
+            cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
+        )
+
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
+
+    # Hoist loop-invariant base offsets
+    base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None]
+    if BLOCK_DPE > 0:
+        base_offs_kpe = cur_kv_head * stride_buf_kh + offs_dpe[:, None]
+    if not HAS_MLA:
+        base_offs_v = cur_kv_head * stride_buf_vh + offs_dv[None, :]
+
+    if split_kv_end > split_kv_start:
+        q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0)
+        q_k = q.to(K_Buffer.dtype.element_ty)
+        if BLOCK_DPE > 0:
+            qpe = tl.load(
+                Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
+            )
+        for start_n in tl.range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+            # Page-aware KV address math (see _fwd_kernel_stage1).
+            if PAGE_SIZE == 1:
+                offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
+            else:
+                page_id = kv_loc // PAGE_SIZE
+                tok_in_p = kv_loc % PAGE_SIZE
+                offs_buf_k = (
+                    page_id[None, :] * stride_buf_kpage
+                    + tok_in_p[None, :] * stride_buf_ktok
+                    + base_offs_k
+                )
+            k = tl.load(
+                K_Buffer + offs_buf_k,
+                mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]),
+                other=0.0,
+            )
+            qk = tl.dot(q_k, k)
+            if BLOCK_DPE > 0:
+                if PAGE_SIZE == 1:
+                    offs_buf_kpe = kv_loc[None, :] * stride_buf_kbs + base_offs_kpe
+                else:
+                    offs_buf_kpe = (
+                        page_id[None, :] * stride_buf_kpage
+                        + tok_in_p[None, :] * stride_buf_ktok
+                        + base_offs_kpe
+                    )
+                kpe = tl.load(
+                    K_Buffer + offs_buf_kpe,
+                    mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]),
+                    other=0.0,
+                )
+                qk += tl.dot(qpe, kpe.to(qpe.dtype))
+            qk *= sm_scale_withk
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            if xai_temperature_len > 0:
+                qk *= xai_temperature_reg[:, None]
+
+            qk = tl.where(
+                mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
+            )
+            if HAS_MLA:
+                v = tl.trans(k)
+            else:
+                if PAGE_SIZE == 1:
+                    offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
+                else:
+                    offs_buf_v = (
+                        page_id[:, None] * stride_buf_vpage
+                        + tok_in_p[:, None] * stride_buf_vtok
+                        + base_offs_v
+                    )
+                v = tl.load(
+                    V_Buffer + offs_buf_v,
+                    mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
+                    other=0.0,
+                )
+
+            n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max[:, None])
+            acc *= re_scale[:, None]
+            acc += tl.dot(p.to(v.dtype), v)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 1)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head[:, None] * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_dv[None, :]
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum[:, None],
+            mask=(mask_h[:, None]) & (mask_dv[None, :]),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+        ) // Lv
+
+        tl.store(
+            Att_Lse + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+            mask=mask_h,
+        )
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def _decode_grouped_att_m_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    att_lse,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    logit_cap,
+    xai_temperature_len=-1,
+    has_mla=False,
+    use_pdl=False,
+    page_size: int = 1,
+):
+    BLOCK = 32
+    Lk = k_buffer.shape[-1]
+    Lv = v_buffer.shape[-1]
+
+    # [TODO] work around shmem limit on MI3xx
+    if _is_hip and Lk >= 576:
+        BLOCK = 16
+
+    if Lk == 576:
+        BLOCK_DMODEL = 512
+        BLOCK_DPE = 64
+    elif Lk == 288:
+        BLOCK_DMODEL = 256
+        BLOCK_DPE = 32
+    else:
+        BLOCK_DMODEL = triton.next_power_of_2(Lk)
+        BLOCK_DPE = 0
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    # 4-D view exposes head_num at dim 2; legacy 3-D exposes
+    # it at dim 1.
+    kv_head_num = k_buffer.shape[-2]
+    batch, head_num = q.shape[0], q.shape[1]
+    kv_group_num = q.shape[1] // kv_head_num
+
+    BLOCK_H = 16
+    MAX_KV_SPLITS = max_kv_splits
+    grid = (
+        batch,
+        triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
+        MAX_KV_SPLITS,
+    )
+
+    extra_kargs = {}
+    num_stages = 2
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
+        num_stages = 1
+
+    k_slot_stride, k_head_stride, k_page_stride, k_tok_stride = _extract_kv_strides(
+        k_buffer, page_size
+    )
+    v_slot_stride, v_head_stride, v_page_stride, v_tok_stride = _extract_kv_strides(
+        v_buffer, page_size
+    )
+
+    _fwd_grouped_kernel_stage1[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        sm_scale_withk,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        att_lse,
+        num_kv_splits,
+        q.stride(0),
+        q.stride(1),
+        k_slot_stride,
+        k_head_stride,
+        v_slot_stride,
+        v_head_stride,
+        k_page_stride,
+        k_tok_stride,
+        v_page_stride,
+        v_tok_stride,
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        kv_group_num=kv_group_num,
+        q_head_num=head_num,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_N=BLOCK,
+        BLOCK_H=BLOCK_H,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        logit_cap=logit_cap,
+        xai_temperature_len=xai_temperature_len,
+        num_warps=4,
+        num_stages=num_stages,
+        Lk=Lk,
+        Lv=Lv,
+        HAS_MLA=has_mla,
+        USE_PDL=use_pdl,
+        PAGE_SIZE=page_size,
+        **extra_kargs,
+    )
+
+
+@triton.jit
+def _fwd_kernel_stage2(
+    Mid_O,
+    Mid_O_1,
+    O,
+    v_scale,
+    kv_indptr,
+    num_kv_splits,
+    sink_ptr,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_obs,
+    stride_oh,
+    MAX_KV_SPLITS: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    Lv: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    USE_PDL: tl.constexpr = False,
+):
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
+    cur_head = tl.program_id(1)
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - tl.load(
+        kv_indptr + cur_batch
+    )
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    offs_d = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lv
+
+    e_sum = 0.0
+    e_max = -float("inf")
+    acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
+
+    offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
+    offs_logic = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh) // Lv
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+
+    for split_kv_id in tl.range(0, MAX_KV_SPLITS, num_stages=2):
+        split_kv_start = kv_len_per_split * split_kv_id
+        split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+        if split_kv_end > split_kv_start:
+            tv = tl.load(
+                Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0
+            )
+            tlogic = tl.load(Mid_O_1 + offs_logic + split_kv_id * stride_mid_os // Lv)
+            n_e_max = tl.maximum(tlogic, e_max)
+
+            old_scale = tl.exp(e_max - n_e_max)
+            acc *= old_scale
+            exp_logic = tl.exp(tlogic - n_e_max)
+            acc += exp_logic * tv
+
+            e_sum = e_sum * old_scale + exp_logic
+            e_max = n_e_max
+
+    if HAS_SINK:
+        cur_sink = tl.load(sink_ptr + cur_head)
+        e_sum += tl.exp(cur_sink - e_max)
+
+    tl.store(
+        O + cur_batch * stride_obs + cur_head * stride_oh + offs_d,
+        acc / e_sum * v_scale,
+        mask=mask_d,
+    )
+
+
+def _decode_softmax_reducev_fwd(
+    logits,
+    lse,
+    q,
+    o,
+    v_scale,
+    v_buffer,
+    kv_indptr,
+    num_kv_splits,
+    max_kv_splits,
+    sinks=None,
+    use_pdl=False,
+):
+    batch, head_num = q.shape[0], q.shape[1]
+    Lv = v_buffer.shape[-1]
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    MAX_KV_SPLITS = max_kv_splits
+    HAS_SINK = sinks is not None
+
+    extra_kargs = {}
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16, "kpack": 2}
+
+    grid = (batch, head_num)
+    _fwd_kernel_stage2[grid](
+        logits,
+        lse,
+        o,
+        v_scale,
+        kv_indptr,
+        num_kv_splits,
+        sinks,
+        logits.stride(0),
+        logits.stride(1),
+        logits.stride(2),
+        o.stride(0),
+        o.stride(1),
+        MAX_KV_SPLITS=MAX_KV_SPLITS,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        BLOCK_DV=BLOCK_DV,
+        Lv=Lv,
+        HAS_SINK=HAS_SINK,
+        USE_PDL=use_pdl,
+        num_warps=4,
+        num_stages=2,
+        **({"launch_pdl": True} if use_pdl else {}),
+        **extra_kargs,
+    )
+
+
+def decode_attention_fwd_normal(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    v_scale,
+    logit_cap=0.0,
+    sinks=None,
+    xai_temperature_len=-1,
+    page_size: int = 1,
+):
+    _decode_att_m_fwd(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        attn_lse,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        max_kv_splits,
+        sm_scale_withk,
+        logit_cap,
+        xai_temperature_len,
+        page_size=page_size,
+    )
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_scale,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+        sinks,
+    )
+
+
+def decode_attention_fwd_grouped(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    v_scale,
+    logit_cap=0.0,
+    sinks=None,
+    xai_temperature_len=-1,
+    has_mla=False,
+    use_pdl=False,
+    page_size: int = 1,
+):
+    _decode_grouped_att_m_fwd(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        attn_lse,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        max_kv_splits,
+        sm_scale_withk,
+        logit_cap,
+        xai_temperature_len,
+        has_mla=has_mla,
+        use_pdl=use_pdl,
+        page_size=page_size,
+    )
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_scale,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+        sinks,
+        use_pdl=use_pdl,
+    )
+
+
+def decode_attention_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    k_scale,
+    v_scale,
+    logit_cap=0.0,
+    sinks=None,
+    xai_temperature_len=-1,
+    has_mla=False,
+    use_pdl=False,
+    page_size: int = 1,
+):
+    assert max_kv_splits == attn_logits.shape[2]
+    assert q.shape[0] <= kv_indptr.shape[0] - 1
+    assert q.shape[0] <= attn_logits.shape[0]
+
+    # head_num lives at dim 1 (3-D) or dim 2 (4-D shared view).
+    kv_head_num = v_buffer.shape[-2]
+    kv_group_num = q.shape[1] // kv_head_num
+
+    if kv_group_num == 1:
+        # MHA
+        decode_attention_fwd_normal(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            attn_lse,
+            num_kv_splits,
+            max_kv_splits,
+            sm_scale * k_scale,
+            v_scale,
+            logit_cap=logit_cap,
+            sinks=sinks,
+            xai_temperature_len=xai_temperature_len,
+            page_size=page_size,
+        )
+    else:
+        # GQA/MQA/MLA
+        decode_attention_fwd_grouped(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            attn_lse,
+            num_kv_splits,
+            max_kv_splits,
+            sm_scale * k_scale,
+            v_scale,
+            logit_cap=logit_cap,
+            sinks=sinks,
+            xai_temperature_len=xai_temperature_len,
+            has_mla=has_mla,
+            use_pdl=use_pdl,
+            page_size=page_size,
+        )

python/sglang/kernels/ops/attention/dsa_metadata.py

@@ -0,0 +1,682 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit(
+    do_not_specialize=[
+        "page_table_stride_0",
+        "real_page_table_stride_0",
+        "max_len",
+    ]
+)
+def _fused_dsa_decode_metadata_kernel(
+    seq_lens,
+    req_pool_indices,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table_1,
+    dsa_cache_seqlens,
+    dsa_cu_seqlens_k,
+    real_page_table,
+    seq_lens_stride: tl.constexpr,
+    req_pool_indices_stride: tl.constexpr,
+    req_to_token_stride_0: tl.constexpr,
+    req_to_token_stride_1: tl.constexpr,
+    page_table_stride_0,
+    page_table_stride_1: tl.constexpr,
+    real_page_table_stride_0,
+    real_page_table_stride_1: tl.constexpr,
+    bs: tl.constexpr,
+    max_len,
+    dsa_index_topk: tl.constexpr,
+    real_page_size: tl.constexpr,
+    HAS_REAL_PAGE_TABLE: tl.constexpr,
+    HAS_PAGE_TABLE_1: tl.constexpr,
+    BLOCK_BS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    if pid == 0:
+        offs_b = tl.arange(0, BLOCK_BS)
+        mask_b = offs_b < bs
+        seq = tl.load(seq_lens + offs_b * seq_lens_stride, mask=mask_b, other=0)
+        seq_i32 = seq.to(tl.int32)
+        dsa_seq = tl.minimum(seq_i32, dsa_index_topk)
+
+        cu = tl.cumsum(seq_i32, 0)
+        dsa_cu = tl.cumsum(dsa_seq, 0)
+
+        tl.store(cache_seqlens + offs_b, seq_i32, mask=mask_b)
+        tl.store(cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(cu_seqlens_k + 1 + offs_b, cu, mask=mask_b)
+        tl.store(dsa_cache_seqlens + offs_b, dsa_seq, mask=mask_b)
+        tl.store(dsa_cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(dsa_cu_seqlens_k + 1 + offs_b, dsa_cu, mask=mask_b)
+        return
+
+    num_col_blocks = tl.cdiv(max_len, BLOCK_N)
+    page_pid = pid - 1
+    row = page_pid // num_col_blocks
+    col_block = page_pid - row * num_col_blocks
+    offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = (row < bs) & (offs_n < max_len)
+
+    req_idx = tl.load(
+        req_pool_indices + row * req_pool_indices_stride,
+        mask=row < bs,
+        other=0,
+    )
+    vals = tl.load(
+        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
+        mask=mask,
+        other=0,
+    ).to(tl.int32)
+    # Write the wide page_size=1 table only when the caller provides it; the
+    # fused decode CUDA graph drops it and consumes real_page_table alone.
+    if HAS_PAGE_TABLE_1:
+        tl.store(
+            page_table_1 + row * page_table_stride_0 + offs_n * page_table_stride_1,
+            vals,
+            mask=mask,
+        )
+
+    if HAS_REAL_PAGE_TABLE:
+        real_mask = mask & ((offs_n % real_page_size) == 0)
+        real_cols = offs_n // real_page_size
+        tl.store(
+            real_page_table
+            + row * real_page_table_stride_0
+            + real_cols * real_page_table_stride_1,
+            vals // real_page_size,
+            mask=real_mask,
+        )
+
+
+def fused_dsa_decode_metadata(
+    seq_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table_1: Optional[torch.Tensor],
+    dsa_cache_seqlens: torch.Tensor,
+    dsa_cu_seqlens_k: torch.Tensor,
+    real_page_table: torch.Tensor,
+    bs: int,
+    max_len: int,
+    dsa_index_topk: int,
+    real_page_size: int,
+) -> None:
+    """Fill decode-graph DSA metadata (seqlens + page tables) from req_to_token.
+
+    ``page_table_1`` (the wide page_size=1 table) is optional: pass ``None`` to
+    skip materializing it and write only the compact ``real_page_table``
+    (page_size=``real_page_size``). This is used by the fused decode CUDA graph,
+    where the wide table is never read (attention uses topk_indices, the indexer
+    uses real_page_table); ``real_page_size`` must be >1 in that case. When a
+    tensor is passed, behavior is unchanged (both tables are written).
+    """
+    assert seq_lens.is_cuda
+    assert req_pool_indices.is_cuda
+    assert req_to_token.is_cuda
+    assert cache_seqlens.is_cuda
+    assert cu_seqlens_k.is_cuda
+    assert dsa_cache_seqlens.is_cuda
+    assert dsa_cu_seqlens_k.is_cuda
+
+    if bs == 0:
+        cu_seqlens_k[:1].zero_()
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+
+    has_real_page_table = real_page_size > 1
+    if has_real_page_table:
+        assert real_page_table is not None
+        assert real_page_table.is_cuda
+    else:
+        # page_size==1: real IS page_table_1, so page_table_1 must be present.
+        assert page_table_1 is not None
+        real_page_table = page_table_1
+
+    # page_table_1 (the wide page_size=1 table) may be dropped for the fused
+    # decode CUDA graph; the kernel then writes only real_page_table.
+    has_page_table_1 = page_table_1 is not None
+    if not has_page_table_1:
+        assert has_real_page_table
+        page_table_1 = real_page_table  # dummy pointer for stride args
+    else:
+        assert page_table_1.is_cuda
+
+    block_bs = triton.next_power_of_2(bs)
+    block_n = 128
+    num_col_blocks = triton.cdiv(max_len, block_n)
+    grid = (1 + bs * num_col_blocks,)
+
+    _fused_dsa_decode_metadata_kernel[grid](
+        seq_lens,
+        req_pool_indices,
+        req_to_token,
+        cache_seqlens,
+        cu_seqlens_k,
+        page_table_1,
+        dsa_cache_seqlens,
+        dsa_cu_seqlens_k,
+        real_page_table,
+        seq_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        req_to_token.stride(1),
+        page_table_1.stride(0),
+        page_table_1.stride(1),
+        real_page_table.stride(0) if has_real_page_table else 0,
+        real_page_table.stride(1) if has_real_page_table else 0,
+        bs,
+        max_len,
+        dsa_index_topk,
+        real_page_size,
+        has_real_page_table,
+        has_page_table_1,
+        BLOCK_BS=block_bs,
+        BLOCK_N=block_n,
+    )
+
+
+@triton.jit(
+    do_not_specialize=[
+        "page_table_stride_0",
+        "real_page_table_stride_0",
+        "max_seqlen_k",
+    ]
+)
+def _fused_dsa_target_verify_metadata_kernel(
+    seq_lens,
+    req_pool_indices,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table_1,
+    seqlens_expanded,
+    dsa_cache_seqlens,
+    dsa_cu_seqlens_k,
+    real_page_table,
+    paged_mqa_ctx_lens_2d,
+    seq_lens_stride: tl.constexpr,
+    req_pool_indices_stride: tl.constexpr,
+    req_to_token_stride_0: tl.constexpr,
+    req_to_token_stride_1: tl.constexpr,
+    page_table_stride_0,
+    page_table_stride_1: tl.constexpr,
+    real_page_table_stride_0,
+    real_page_table_stride_1: tl.constexpr,
+    paged_mqa_ctx_lens_stride_0: tl.constexpr,
+    paged_mqa_ctx_lens_stride_1: tl.constexpr,
+    bs: tl.constexpr,
+    max_seqlen_k,
+    dsa_index_topk: tl.constexpr,
+    real_page_size: tl.constexpr,
+    next_n: tl.constexpr,
+    HAS_REAL_PAGE_TABLE: tl.constexpr,
+    HAS_PAGED_MQA_CTX_LENS: tl.constexpr,
+    HAS_PAGE_TABLE_1: tl.constexpr,
+    BLOCK_BS: tl.constexpr,
+    BLOCK_EXPANDED: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    expanded_size: tl.constexpr = bs * next_n
+
+    if pid == 0:
+        offs_b = tl.arange(0, BLOCK_BS)
+        mask_b = offs_b < bs
+        seq = tl.load(seq_lens + offs_b * seq_lens_stride, mask=mask_b, other=0)
+        cache_seq = seq.to(tl.int32) + next_n
+        cu = tl.cumsum(cache_seq, 0)
+
+        tl.store(cache_seqlens + offs_b, cache_seq, mask=mask_b)
+        tl.store(cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(cu_seqlens_k + 1 + offs_b, cu, mask=mask_b)
+
+        offs_e = tl.arange(0, BLOCK_EXPANDED)
+        mask_e = offs_e < expanded_size
+        req_row = offs_e // next_n
+        draft_off = offs_e - req_row * next_n
+        base_seq = tl.load(
+            seq_lens + req_row * seq_lens_stride,
+            mask=mask_e,
+            other=0,
+        ).to(tl.int32)
+        expanded_seq = base_seq + draft_off + 1
+        expanded_seq = tl.where(mask_e, expanded_seq, 0)
+        dsa_seq = tl.minimum(expanded_seq, dsa_index_topk)
+        dsa_cu = tl.cumsum(dsa_seq, 0)
+
+        tl.store(seqlens_expanded + offs_e, expanded_seq, mask=mask_e)
+        tl.store(dsa_cache_seqlens + offs_e, dsa_seq, mask=mask_e)
+        tl.store(dsa_cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(dsa_cu_seqlens_k + 1 + offs_e, dsa_cu, mask=mask_e)
+
+        if HAS_PAGED_MQA_CTX_LENS:
+            tl.store(
+                paged_mqa_ctx_lens_2d
+                + req_row * paged_mqa_ctx_lens_stride_0
+                + draft_off * paged_mqa_ctx_lens_stride_1,
+                base_seq + next_n,
+                mask=mask_e,
+            )
+        return
+
+    num_col_blocks = tl.cdiv(max_seqlen_k, BLOCK_N)
+    page_pid = pid - 1
+    out_row = page_pid // num_col_blocks
+    col_block = page_pid - out_row * num_col_blocks
+    offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = (out_row < expanded_size) & (offs_n < max_seqlen_k)
+
+    req_row = out_row // next_n
+    req_idx = tl.load(
+        req_pool_indices + req_row * req_pool_indices_stride,
+        mask=out_row < expanded_size,
+        other=0,
+    )
+    vals = tl.load(
+        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
+        mask=mask,
+        other=0,
+    ).to(tl.int32)
+    # Write the wide page_size=1 table only when the caller provides it (see
+    # fused_dsa_decode_metadata for the optional-page_table_1 contract).
+    if HAS_PAGE_TABLE_1:
+        tl.store(
+            page_table_1 + out_row * page_table_stride_0 + offs_n * page_table_stride_1,
+            vals,
+            mask=mask,
+        )
+
+    if HAS_REAL_PAGE_TABLE:
+        real_mask = mask & ((offs_n % real_page_size) == 0)
+        real_cols = offs_n // real_page_size
+        tl.store(
+            real_page_table
+            + out_row * real_page_table_stride_0
+            + real_cols * real_page_table_stride_1,
+            vals // real_page_size,
+            mask=real_mask,
+        )
+
+
+def fused_dsa_target_verify_metadata(
+    seq_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table_1: Optional[torch.Tensor],
+    seqlens_expanded: torch.Tensor,
+    dsa_cache_seqlens: torch.Tensor,
+    dsa_cu_seqlens_k: torch.Tensor,
+    real_page_table: torch.Tensor,
+    bs: int,
+    max_seqlen_k: int,
+    dsa_index_topk: int,
+    real_page_size: int,
+    next_n: int,
+    paged_mqa_ctx_lens_2d: torch.Tensor = None,
+) -> None:
+    assert seq_lens.is_cuda
+    assert req_pool_indices.is_cuda
+    assert req_to_token.is_cuda
+    assert cache_seqlens.is_cuda
+    assert cu_seqlens_k.is_cuda
+    assert seqlens_expanded.is_cuda
+    assert dsa_cache_seqlens.is_cuda
+    assert dsa_cu_seqlens_k.is_cuda
+
+    if bs == 0:
+        cu_seqlens_k[:1].zero_()
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+    assert next_n > 0
+
+    has_real_page_table = real_page_size > 1
+    if has_real_page_table:
+        assert real_page_table is not None
+        assert real_page_table.is_cuda
+    else:
+        assert page_table_1 is not None
+        real_page_table = page_table_1
+
+    # page_table_1 (the wide page_size=1 table) may be dropped for the fused
+    # decode CUDA graph; the kernel then writes only real_page_table.
+    has_page_table_1 = page_table_1 is not None
+    if not has_page_table_1:
+        assert has_real_page_table
+        page_table_1 = real_page_table  # dummy pointer for stride args
+    else:
+        assert page_table_1.is_cuda
+
+    has_paged_mqa_ctx_lens = paged_mqa_ctx_lens_2d is not None
+    if has_paged_mqa_ctx_lens:
+        assert paged_mqa_ctx_lens_2d.is_cuda
+        assert paged_mqa_ctx_lens_2d.dtype == torch.int32
+        assert paged_mqa_ctx_lens_2d.dim() == 2
+        assert paged_mqa_ctx_lens_2d.size(0) == bs
+        assert paged_mqa_ctx_lens_2d.size(1) == next_n
+    else:
+        paged_mqa_ctx_lens_2d = page_table_1
+
+    expanded_size = bs * next_n
+    block_bs = triton.next_power_of_2(bs)
+    block_expanded = triton.next_power_of_2(expanded_size)
+    block_n = 128
+    num_col_blocks = triton.cdiv(max_seqlen_k, block_n)
+    grid = (1 + expanded_size * num_col_blocks,)
+
+    _fused_dsa_target_verify_metadata_kernel[grid](
+        seq_lens,
+        req_pool_indices,
+        req_to_token,
+        cache_seqlens,
+        cu_seqlens_k,
+        page_table_1,
+        seqlens_expanded,
+        dsa_cache_seqlens,
+        dsa_cu_seqlens_k,
+        real_page_table,
+        paged_mqa_ctx_lens_2d,
+        seq_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        req_to_token.stride(1),
+        page_table_1.stride(0),
+        page_table_1.stride(1),
+        real_page_table.stride(0) if has_real_page_table else 0,
+        real_page_table.stride(1) if has_real_page_table else 0,
+        paged_mqa_ctx_lens_2d.stride(0) if has_paged_mqa_ctx_lens else 0,
+        paged_mqa_ctx_lens_2d.stride(1) if has_paged_mqa_ctx_lens else 0,
+        bs,
+        max_seqlen_k,
+        dsa_index_topk,
+        real_page_size,
+        next_n,
+        has_real_page_table,
+        has_paged_mqa_ctx_lens,
+        has_page_table_1,
+        BLOCK_BS=block_bs,
+        BLOCK_EXPANDED=block_expanded,
+        BLOCK_N=block_n,
+    )
+
+
+@triton.jit(
+    do_not_specialize=[
+        "page_table_stride_0",
+        "real_page_table_stride_0",
+        "total_len",
+        "max_seqlen_k",
+    ]
+)
+def _fused_dsa_draft_extend_metadata_kernel(
+    seq_lens,
+    extend_seq_lens,
+    req_pool_indices,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table_1,
+    seqlens_expanded,
+    dsa_cache_seqlens,
+    dsa_cu_seqlens_k,
+    real_page_table,
+    seq_lens_stride: tl.constexpr,
+    extend_seq_lens_stride: tl.constexpr,
+    req_pool_indices_stride: tl.constexpr,
+    req_to_token_stride_0: tl.constexpr,
+    req_to_token_stride_1: tl.constexpr,
+    page_table_stride_0,
+    page_table_stride_1: tl.constexpr,
+    real_page_table_stride_0,
+    real_page_table_stride_1: tl.constexpr,
+    bs: tl.constexpr,
+    total_len,
+    max_seqlen_k,
+    dsa_index_topk: tl.constexpr,
+    real_page_size: tl.constexpr,
+    HAS_REAL_PAGE_TABLE: tl.constexpr,
+    HAS_PAGE_TABLE_1: tl.constexpr,
+    STATIC_EXTEND_LEN: tl.constexpr,
+    BLOCK_BS: tl.constexpr,
+    BLOCK_EXPANDED: tl.constexpr,
+    BLOCK_ROWS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    if pid == 0:
+        offs_b = tl.arange(0, BLOCK_BS)
+        mask_b = offs_b < bs
+        seq = tl.load(seq_lens + offs_b * seq_lens_stride, mask=mask_b, other=0)
+        cache_seq = seq.to(tl.int32)
+        cu = tl.cumsum(cache_seq, 0)
+
+        tl.store(cache_seqlens + offs_b, cache_seq, mask=mask_b)
+        tl.store(cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(cu_seqlens_k + 1 + offs_b, cu, mask=mask_b)
+
+        offs_e = tl.arange(0, BLOCK_EXPANDED)
+        mask_e = offs_e < total_len
+        if STATIC_EXTEND_LEN:
+            static_qo_len = tl.load(extend_seq_lens).to(tl.int32)
+            req_row = offs_e // static_qo_len
+            local_off = offs_e - req_row * static_qo_len
+            qo_len_for_row = tl.zeros((BLOCK_EXPANDED,), tl.int32) + static_qo_len
+        else:
+            req_row = tl.full((BLOCK_EXPANDED,), 0, tl.int32)
+            local_off = tl.full((BLOCK_EXPANDED,), 0, tl.int32)
+            qo_len_for_row = tl.full((BLOCK_EXPANDED,), 1, tl.int32)
+            prefix = tl.full((), 0, tl.int32)
+
+            for i in tl.range(0, bs):
+                qo_len = tl.load(extend_seq_lens + i * extend_seq_lens_stride).to(
+                    tl.int32
+                )
+                in_row = (offs_e >= prefix) & (offs_e < prefix + qo_len)
+                req_row = tl.where(in_row, i, req_row)
+                local_off = tl.where(in_row, offs_e - prefix, local_off)
+                qo_len_for_row = tl.where(in_row, qo_len, qo_len_for_row)
+                prefix += qo_len
+
+        base_seq = tl.load(
+            seq_lens + req_row * seq_lens_stride,
+            mask=mask_e,
+            other=0,
+        ).to(tl.int32)
+        # Clamp to >= 0: DP-padded / idle-companion rows carry the CUDA-graph
+        # seq_len fill value (1), which is smaller than qo_len, so the raw
+        # per-row visible kv length goes negative. Consumers treat these
+        # lengths as unsigned (the top-k v2 kernel reads them as uint32), so a
+        # negative row becomes a ~4e9-token length and an illegal memory
+        # access. 0 keeps padded rows on the trivial all-(-1) output path.
+        expanded_seq = base_seq - qo_len_for_row + local_off + 1
+        expanded_seq = tl.maximum(expanded_seq, 0)
+        expanded_seq = tl.where(mask_e, expanded_seq, 0)
+        dsa_seq = tl.minimum(expanded_seq, dsa_index_topk)
+        dsa_cu = tl.cumsum(dsa_seq, 0)
+
+        tl.store(seqlens_expanded + offs_e, expanded_seq, mask=mask_e)
+        tl.store(dsa_cache_seqlens + offs_e, dsa_seq, mask=mask_e)
+        tl.store(dsa_cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(dsa_cu_seqlens_k + 1 + offs_e, dsa_cu, mask=mask_e)
+        return
+
+    num_col_blocks = tl.cdiv(max_seqlen_k, BLOCK_N)
+    page_pid = pid - 1
+    req_row = page_pid // num_col_blocks
+    col_block = page_pid - req_row * num_col_blocks
+    offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    qo_len = tl.load(
+        extend_seq_lens + req_row * extend_seq_lens_stride,
+        mask=req_row < bs,
+        other=0,
+    ).to(tl.int32)
+    if STATIC_EXTEND_LEN:
+        prefix = req_row * qo_len
+    else:
+        prefix = tl.full((), 0, tl.int32)
+        for i in tl.range(0, bs):
+            prev_qo_len = tl.load(extend_seq_lens + i * extend_seq_lens_stride).to(
+                tl.int32
+            )
+            prefix += tl.where(i < req_row, prev_qo_len, 0)
+    offs_r = tl.arange(0, BLOCK_ROWS)
+    out_rows = prefix + offs_r
+    row_mask = (req_row < bs) & (offs_r < qo_len) & (out_rows < total_len)
+    col_mask = offs_n < max_seqlen_k
+    has_rows = (req_row < bs) & (qo_len > 0)
+    mask = row_mask[:, None] & col_mask[None, :]
+
+    req_idx = tl.load(
+        req_pool_indices + req_row * req_pool_indices_stride,
+        mask=has_rows,
+        other=0,
+    )
+    vals = tl.load(
+        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
+        mask=col_mask & has_rows,
+        other=0,
+    ).to(tl.int32)
+    # Write the wide page_size=1 table only when the caller provides it (see
+    # fused_dsa_decode_metadata for the optional-page_table_1 contract).
+    if HAS_PAGE_TABLE_1:
+        tl.store(
+            page_table_1
+            + out_rows[:, None] * page_table_stride_0
+            + offs_n[None, :] * page_table_stride_1,
+            vals[None, :],
+            mask=mask,
+        )
+
+    if HAS_REAL_PAGE_TABLE:
+        real_mask = mask & ((offs_n[None, :] % real_page_size) == 0)
+        real_cols = offs_n // real_page_size
+        tl.store(
+            real_page_table
+            + out_rows[:, None] * real_page_table_stride_0
+            + real_cols[None, :] * real_page_table_stride_1,
+            (vals // real_page_size)[None, :],
+            mask=real_mask,
+        )
+
+
+def fused_dsa_draft_extend_metadata(
+    seq_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table_1: Optional[torch.Tensor],
+    seqlens_expanded: torch.Tensor,
+    dsa_cache_seqlens: torch.Tensor,
+    dsa_cu_seqlens_k: torch.Tensor,
+    real_page_table: torch.Tensor,
+    bs: int,
+    total_len: int,
+    max_seqlen_k: int,
+    dsa_index_topk: int,
+    real_page_size: int,
+    max_extend_len: int,
+    max_total_len: int,
+    static_extend_len: bool = False,
+) -> None:
+    assert seq_lens.is_cuda
+    assert extend_seq_lens.is_cuda
+    assert req_pool_indices.is_cuda
+    assert req_to_token.is_cuda
+    assert cache_seqlens.is_cuda
+    assert cu_seqlens_k.is_cuda
+    assert seqlens_expanded.is_cuda
+    assert dsa_cache_seqlens.is_cuda
+    assert dsa_cu_seqlens_k.is_cuda
+
+    if bs == 0:
+        cu_seqlens_k[:1].zero_()
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+    if total_len == 0:
+        cache = seq_lens.to(torch.int32)
+        cache_seqlens.copy_(cache)
+        cu_seqlens_k[:1].zero_()
+        cu_seqlens_k[1 : bs + 1].copy_(torch.cumsum(cache, dim=0, dtype=torch.int32))
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+    assert total_len <= max_total_len
+    # Caller-owned graph metadata guarantees each request accepts at most
+    # max_extend_len tokens. Avoid checking extend_seq_lens.max() here because
+    # that would sync in the replay hot path.
+    assert max_extend_len > 0
+    assert total_len <= bs * max_extend_len
+
+    has_real_page_table = real_page_size > 1
+    if has_real_page_table:
+        assert real_page_table is not None
+        assert real_page_table.is_cuda
+    else:
+        assert page_table_1 is not None
+        real_page_table = page_table_1
+
+    # page_table_1 (the wide page_size=1 table) may be dropped for the fused
+    # decode CUDA graph; the kernel then writes only real_page_table.
+    has_page_table_1 = page_table_1 is not None
+    if not has_page_table_1:
+        assert has_real_page_table
+        page_table_1 = real_page_table  # dummy pointer for stride args
+    else:
+        assert page_table_1.is_cuda
+
+    block_bs = triton.next_power_of_2(bs)
+    block_expanded = triton.next_power_of_2(max_total_len)
+    block_rows = triton.next_power_of_2(max_extend_len)
+    block_n = 128
+    num_col_blocks = triton.cdiv(max_seqlen_k, block_n)
+    grid = (1 + bs * num_col_blocks,)
+
+    _fused_dsa_draft_extend_metadata_kernel[grid](
+        seq_lens,
+        extend_seq_lens,
+        req_pool_indices,
+        req_to_token,
+        cache_seqlens,
+        cu_seqlens_k,
+        page_table_1,
+        seqlens_expanded,
+        dsa_cache_seqlens,
+        dsa_cu_seqlens_k,
+        real_page_table,
+        seq_lens.stride(0),
+        extend_seq_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        req_to_token.stride(1),
+        page_table_1.stride(0),
+        page_table_1.stride(1),
+        real_page_table.stride(0) if has_real_page_table else 0,
+        real_page_table.stride(1) if has_real_page_table else 0,
+        bs,
+        total_len,
+        max_seqlen_k,
+        dsa_index_topk,
+        real_page_size,
+        has_real_page_table,
+        has_page_table_1,
+        static_extend_len,
+        BLOCK_BS=block_bs,
+        BLOCK_EXPANDED=block_expanded,
+        BLOCK_ROWS=block_rows,
+        BLOCK_N=block_n,
+    )

python/sglang/kernels/ops/attention/merge_state.py

@@ -0,0 +1,96 @@
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def merge_state_kernel(
+    output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_merged
+    output_lse,  # [NUM_TOKENS, NUM_HEADS] s_merged
+    prefix_output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_a
+    prefix_lse,  # [NUM_TOKENS, NUM_HEADS] s_a
+    suffix_output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_b
+    suffix_lse,  # [NUM_TOKENS, NUM_HEADS] s_b
+    HEAD_SIZE: tl.constexpr,
+    PADDED_HEAD_SIZE: tl.constexpr,
+    OUTPUT_LSE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    num_tokens = tl.num_programs(0)
+    head_idx = tl.program_id(1)
+    num_heads = tl.num_programs(1)
+
+    p_lse = tl.load(prefix_lse + token_idx * num_heads + head_idx)
+    s_lse = tl.load(suffix_lse + token_idx * num_heads + head_idx)
+    p_lse = float("-inf") if p_lse == float("inf") else p_lse
+    s_lse = float("-inf") if s_lse == float("inf") else s_lse
+
+    max_lse = tl.maximum(p_lse, s_lse)
+    p_lse = p_lse - max_lse
+    s_lse = s_lse - max_lse
+    out_se = tl.exp(p_lse) + tl.exp(s_lse)
+
+    if OUTPUT_LSE:
+        out_lse = tl.log(out_se) + max_lse
+        tl.store(output_lse + token_idx * num_heads + head_idx, out_lse)
+
+    head_arange = tl.arange(0, PADDED_HEAD_SIZE)
+    head_mask = head_arange < HEAD_SIZE
+    p_out = tl.load(
+        prefix_output
+        + token_idx * num_heads * HEAD_SIZE
+        + head_idx * HEAD_SIZE
+        + head_arange,
+        mask=head_mask,
+    )
+    s_out = tl.load(
+        suffix_output
+        + token_idx * num_heads * HEAD_SIZE
+        + head_idx * HEAD_SIZE
+        + head_arange,
+        mask=head_mask,
+    )
+
+    p_scale = tl.exp(p_lse) / out_se
+    s_scale = tl.exp(s_lse) / out_se
+    out = p_out * p_scale + s_out * s_scale
+    tl.store(
+        output + token_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE + head_arange,
+        out,
+        mask=head_mask,
+    )
+
+
+def merge_state_triton(
+    prefix_output: torch.Tensor,
+    prefix_lse: torch.Tensor,
+    suffix_output: torch.Tensor,
+    suffix_lse: torch.Tensor,
+    output: Optional[torch.Tensor] = None,
+    output_lse: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    # Avoid creating new tensors if they are already provided
+    if output is None:
+        output = torch.empty_like(prefix_output)
+    if output_lse is None:
+        output_lse = torch.empty_like(prefix_lse)
+
+    num_tokens = output.shape[0]
+    num_query_heads = output.shape[1]
+    head_size = output.shape[2]
+    padded_head_size = triton.next_power_of_2(head_size)
+
+    merge_state_kernel[(num_tokens, num_query_heads)](
+        output,
+        output_lse,
+        prefix_output,
+        prefix_lse,
+        suffix_output,
+        suffix_lse,
+        head_size,
+        padded_head_size,
+        output_lse is not None,
+    )
+    return output, output_lse

python/sglang/kernels/ops/attention/metadata.py

@@ -0,0 +1,467 @@
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import triton
+import triton.language as tl
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+
+
+@triton.jit
+def get_num_kv_splits_triton(
+    num_kv_splits_ptr,
+    seq_lens_ptr,
+    num_seq,
+    num_group,
+    num_head,
+    num_kv_head,
+    max_kv_splits,
+    device_core_count,
+    MAX_NUM_SEQ: tl.constexpr,
+):
+    # TODO: this method is tunable, we need more online serving data to tune it
+    offs_seq = tl.arange(0, MAX_NUM_SEQ)
+    mask_seq = offs_seq < num_seq
+
+    seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0)
+    max_seq_len = tl.max(seq_lens)
+    seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len)
+    min_seq_len = tl.min(seq_lens)
+    if max_seq_len * 8 < min_seq_len * 10:
+        min_seq_len = max_seq_len
+    max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits)
+    kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1)
+
+    # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually
+    ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0
+    ext_device_core_count = tl.cast(
+        device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32
+    )
+    block_h, num_kv_group = 16, num_head // num_kv_head
+    if num_kv_group == 1:
+        token_grid = num_seq * num_group * num_head
+    else:
+        # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd
+        block_h = tl.minimum(block_h, num_kv_group)
+        token_grid = num_seq * num_group * tl.cdiv(num_head, block_h)
+    max_kv_splits_2 = tl.minimum(
+        tl.cdiv(ext_device_core_count, token_grid), max_kv_splits
+    )
+    kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2)
+
+    num_kv_splits = tl.maximum(
+        tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2)
+    )
+
+    offs_token = offs_seq * num_group
+    mask_token = offs_token < num_seq * num_group
+    for i in range(0, num_group):
+        tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token)
+
+
+@triton.jit
+def _prepare_swa_spec_page_table_kernel(
+    dst_ptr,
+    src_a_ptr,
+    src_b_ptr,
+    seq_len_a_ptr,
+    seq_len_b_ptr,
+    dst_stride_m,
+    dst_stride_n,
+    a_stride_m,
+    a_stride_n,
+    b_stride_m,
+    b_stride_n,
+    LEN_A: tl.constexpr,
+    LEN_B: tl.constexpr,
+    REPEAT_STEP: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    idx_a = pid_m // REPEAT_STEP
+    idx_b = pid_m
+    seq_len_a = tl.load(seq_len_a_ptr + idx_a)
+    seq_len_b = tl.load(seq_len_b_ptr + idx_b)
+
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    total_len = seq_len_a + seq_len_b
+
+    if pid_n * BLOCK_N >= total_len:
+        return
+
+    mask = offs_n < total_len
+    dst = dst_ptr + pid_m * dst_stride_m + offs_n * dst_stride_n
+
+    if (pid_n + 1) * BLOCK_N < seq_len_a:
+        a_ptr = src_a_ptr + idx_a * a_stride_m + offs_n * a_stride_n
+        a_mask = mask & (offs_n < LEN_A)
+        val = tl.load(a_ptr, mask=a_mask, other=0)
+        tl.store(dst, val, mask=mask)
+    elif pid_n * BLOCK_N >= seq_len_a:
+        offs_b = offs_n - seq_len_a
+        b_ptr = src_b_ptr + idx_b * b_stride_m + offs_b * b_stride_n
+        b_mask = mask & (offs_b < LEN_B)
+        val = tl.load(b_ptr, mask=b_mask, other=0)
+        tl.store(dst, val, mask=mask)
+    else:
+        # mixed part
+        a_offs = offs_n
+        a_mask = (a_offs < seq_len_a) & (a_offs < LEN_A)
+        a_ptr = src_a_ptr + idx_a * a_stride_m + a_offs * a_stride_n
+        a_val = tl.load(a_ptr, mask=a_mask, other=0)
+
+        b_offs = offs_n - seq_len_a
+        b_mask = (b_offs >= 0) & (b_offs < seq_len_b) & (b_offs < LEN_B)
+        b_ptr = src_b_ptr + idx_b * b_stride_m + b_offs * b_stride_n
+        b_val = tl.load(b_ptr, mask=b_mask, other=0)
+
+        result = tl.where(offs_n < seq_len_a, a_val, b_val)
+        tl.store(dst, result, mask=mask)
+
+
+def prepare_swa_spec_page_table_triton(
+    page_table_dst: torch.Tensor,
+    page_table_a: torch.Tensor,
+    page_table_b: torch.Tensor,  # expand page table
+    seq_len_a: torch.Tensor,
+    seq_len_b: torch.Tensor,  # expand seq lens
+    speculative_num_draft_tokens: int,
+):
+    # concat page_table and expand page_table by kv seq length
+    bs = seq_len_a.numel()
+    bs_expand = seq_len_b.numel()
+    assert bs_expand == bs * speculative_num_draft_tokens
+
+    LEN_A = page_table_a.shape[1]
+    LEN_B = page_table_b.shape[1]
+    LEN_OUT = LEN_A + LEN_B
+    REPEAT_STEP = speculative_num_draft_tokens
+    BLOCK_N = 256
+
+    grid = (bs_expand, triton.cdiv(LEN_OUT, BLOCK_N))
+    _prepare_swa_spec_page_table_kernel[grid](
+        page_table_dst,
+        page_table_a,
+        page_table_b,
+        seq_len_a,
+        seq_len_b,
+        page_table_dst.stride(0),
+        page_table_dst.stride(1),
+        page_table_a.stride(0),
+        page_table_a.stride(1),
+        page_table_b.stride(0),
+        page_table_b.stride(1),
+        LEN_A=LEN_A,
+        LEN_B=LEN_B,
+        REPEAT_STEP=REPEAT_STEP,
+        BLOCK_N=BLOCK_N,
+        num_warps=4,
+    )
+
+
+@triton.jit
+def _fused_metadata_kernel_general(
+    # Input tensors
+    seq_lens,
+    seq_lens_stride_0,
+    req_to_token,
+    req_to_token_stride_0,
+    req_to_token_stride_1,
+    req_pool_indices,
+    req_pool_indices_stride_0,
+    # Output buffers
+    cache_seqlens_int32,
+    cache_seqlens_int32_stride_0,
+    cu_seqlens_k,
+    cu_seqlens_k_stride_0,
+    page_table,
+    page_table_stride_0,
+    page_table_stride_1,
+    swa_page_table,
+    swa_page_table_stride_0,
+    swa_page_table_stride_1,
+    full_to_swa_mapping,
+    full_to_swa_mapping_stride_0,
+    # Scalar parameters
+    B,
+    max_seq_pages,
+    page_size: tl.constexpr,
+    seq_len_delta: tl.constexpr,
+    use_swa: tl.constexpr,
+    SHIFT: tl.constexpr,
+    BLOCK_COLS: tl.constexpr,
+):
+    pid_b = tl.program_id(0)  # batch index
+    pid_c = tl.program_id(1)  # column chunk index
+
+    # 1. Prefix sum (only one block does it)
+    if pid_b == 0 and pid_c == 0:
+        acc = 0
+        for idx in range(B):
+            seq = tl.load(seq_lens + idx * seq_lens_stride_0)
+            val = (seq + seq_len_delta).to(tl.int32)
+            tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val)
+            tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc)
+            acc += val
+        tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc)
+
+    # 2. Gather for this batch and column chunk
+    if max_seq_pages == 0:
+        return
+
+    i = pid_b
+    # Load row index for this batch (all threads in block have same i)
+    row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0)
+    row_offset = row_idx * req_to_token_stride_0
+
+    col_start = pid_c * BLOCK_COLS
+    col_offsets = col_start + tl.arange(0, BLOCK_COLS)
+    mask = col_offsets < max_seq_pages
+
+    # Compute column indices in the source tensor (token offset)
+    if page_size == 1:
+        col_idx = col_offsets
+    else:
+        col_idx = col_offsets << SHIFT  # faster than multiplication for power-of-two
+
+    # Load page indices from req_to_token
+    rt_offsets = row_offset + col_idx * req_to_token_stride_1
+    page_index = tl.load(
+        req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg"
+    )
+
+    # Compute page_table
+    if page_size == 1:
+        page_table_val = page_index
+    else:
+        page_table_val = page_index >> SHIFT
+
+    # Store to page_table
+    pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
+    tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg")
+
+    if use_swa:
+        swa_slot = tl.load(
+            full_to_swa_mapping + page_index * full_to_swa_mapping_stride_0,
+            mask=mask,
+            other=0,
+            cache_modifier=".cg",
+        )
+        if page_size == 1:
+            swa_val = swa_slot
+        else:
+            swa_val = swa_slot >> SHIFT
+        swa_offsets = (
+            i * swa_page_table_stride_0 + col_offsets * swa_page_table_stride_1
+        )
+        tl.store(swa_page_table + swa_offsets, swa_val, mask=mask, cache_modifier=".cg")
+
+
+@triton.jit
+def _fused_metadata_kernel_ps1_no_swa(
+    # Input tensors
+    seq_lens,
+    seq_lens_stride_0,
+    req_to_token,
+    req_to_token_stride_0,
+    req_to_token_stride_1,
+    req_pool_indices,
+    req_pool_indices_stride_0,
+    # Output buffers
+    cache_seqlens_int32,
+    cache_seqlens_int32_stride_0,
+    cu_seqlens_k,
+    cu_seqlens_k_stride_0,
+    page_table,
+    page_table_stride_0,
+    page_table_stride_1,
+    # Scalar parameters
+    B,
+    max_seq_pages,
+    seq_len_delta: tl.constexpr,
+    BLOCK_COLS: tl.constexpr,
+):
+    pid_b = tl.program_id(0)  # batch index
+    pid_c = tl.program_id(1)  # column chunk index
+
+    # 1. Prefix sum (only one block does it)
+    if pid_b == 0 and pid_c == 0:
+        acc = 0
+        for idx in range(B):
+            seq = tl.load(seq_lens + idx * seq_lens_stride_0)
+            val = (seq + seq_len_delta).to(tl.int32)
+            tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val)
+            tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc)
+            acc += val
+        tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc)
+
+    # 2. Gather for this batch and column chunk
+    if max_seq_pages == 0:
+        return
+
+    i = pid_b
+    # Load row index for this batch (all threads in block have same i)
+    row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0)
+    row_offset = row_idx * req_to_token_stride_0
+
+    col_start = pid_c * BLOCK_COLS
+    col_offsets = col_start + tl.arange(0, BLOCK_COLS)
+    mask = col_offsets < max_seq_pages
+
+    # page_size = 1: col_idx = col_offsets
+    rt_offsets = row_offset + col_offsets * req_to_token_stride_1
+    page_index = tl.load(
+        req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg"
+    )
+
+    # page_table = page_index // 1 = page_index
+    pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
+    tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg")
+
+
+def normal_decode_set_metadata(
+    cache_seqlens_int32: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table: torch.Tensor,
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    strided_indices: torch.Tensor,
+    max_seq_pages: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seq_len_delta: int,
+    page_size: int,
+    swa_page_table: Optional[torch.Tensor] = None,
+    token_to_kv_pool: Optional["SWAKVPool"] = None,
+):
+    """
+    Fused Triton implementation that replaces 4-5 sequential CUDA kernels with 1-2 kernels:
+      1. cache_seqlens = seq_lens + seq_len_delta (int64->int32 cast)
+      2. cu_seqlens_k = cumsum(cache_seqlens) (prefix-sum)
+      3. page_indices = req_to_token[pool_idx, stride_idx] (2-D gather)
+      4. page_table = page_indices // page_size (floor-divide)
+      5. (optional) swa_page_table for sliding window attention
+
+    Achieves ~5.2x speedup on H200 hardware for typical decode workloads.
+    """
+    assert (
+        page_size > 0 and (page_size & (page_size - 1)) == 0
+    ), f"page_size must be a power of two, got {page_size}"
+
+    batch_size = cache_seqlens_int32.shape[0]
+    device = seq_lens.device
+
+    # Ensure contiguous memory layout for efficient Triton access
+    seq_lens = seq_lens.contiguous()
+    req_to_token = req_to_token.contiguous()
+    req_pool_indices = req_pool_indices.contiguous()
+
+    # Prepare tensor strides
+    seq_lens_stride_0 = seq_lens.stride(0)
+    req_to_token_stride_0 = req_to_token.stride(0)
+    req_to_token_stride_1 = req_to_token.stride(1)
+    req_pool_indices_stride_0 = req_pool_indices.stride(0)
+    cache_seqlens_int32_stride_0 = cache_seqlens_int32.stride(0)
+    cu_seqlens_k_stride_0 = cu_seqlens_k.stride(0)
+    page_table_stride_0 = page_table.stride(0)
+    page_table_stride_1 = page_table.stride(1)
+
+    # Check if we should use the specialized fast path for page_size=1, no SWA
+    use_swa = swa_page_table is not None and token_to_kv_pool is not None
+
+    if page_size == 1 and not use_swa:
+        # Specialized kernel for the common case (page_size=1, no SWA)
+        BLOCK_COLS = 256
+        if max_seq_pages == 0:
+            grid = (1, 1)
+        else:
+            num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS)
+            grid = (batch_size, num_blocks_j)
+
+        _fused_metadata_kernel_ps1_no_swa[grid](
+            seq_lens,
+            seq_lens_stride_0,
+            req_to_token,
+            req_to_token_stride_0,
+            req_to_token_stride_1,
+            req_pool_indices,
+            req_pool_indices_stride_0,
+            cache_seqlens_int32,
+            cache_seqlens_int32_stride_0,
+            cu_seqlens_k,
+            cu_seqlens_k_stride_0,
+            page_table,
+            page_table_stride_0,
+            page_table_stride_1,
+            batch_size,
+            max_seq_pages,
+            seq_len_delta,
+            BLOCK_COLS=BLOCK_COLS,
+            num_warps=8,
+            num_stages=3,
+        )
+    else:
+        # General kernel for page_size > 1 or SWA cases
+        # SWA parameters
+        if use_swa:
+            from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+
+            assert isinstance(token_to_kv_pool, SWAKVPool)
+            swa_page_table = swa_page_table.contiguous()
+            swa_page_table_stride_0 = swa_page_table.stride(0)
+            swa_page_table_stride_1 = swa_page_table.stride(1)
+            # Extract the full_to_swa_index_mapping from token_to_kv_pool
+            full_to_swa_mapping = (
+                token_to_kv_pool.full_to_swa_index_mapping.contiguous()
+            )
+            full_to_swa_mapping_stride_0 = full_to_swa_mapping.stride(0)
+        else:
+            # Dummy tensors (not used)
+            swa_page_table = torch.empty(0, dtype=torch.int32, device=device)
+            swa_page_table_stride_0 = 0
+            swa_page_table_stride_1 = 0
+            full_to_swa_mapping = torch.empty(0, dtype=torch.int32, device=device)
+            full_to_swa_mapping_stride_0 = 0
+
+        # Kernel configuration
+        BLOCK_COLS = 128
+        shift = (page_size).bit_length() - 1 if page_size > 1 else 0
+
+        if max_seq_pages == 0:
+            grid = (1, 1)
+        else:
+            num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS)
+            grid = (batch_size, num_blocks_j)
+
+        _fused_metadata_kernel_general[grid](
+            seq_lens,
+            seq_lens_stride_0,
+            req_to_token,
+            req_to_token_stride_0,
+            req_to_token_stride_1,
+            req_pool_indices,
+            req_pool_indices_stride_0,
+            cache_seqlens_int32,
+            cache_seqlens_int32_stride_0,
+            cu_seqlens_k,
+            cu_seqlens_k_stride_0,
+            page_table,
+            page_table_stride_0,
+            page_table_stride_1,
+            swa_page_table,
+            swa_page_table_stride_0,
+            swa_page_table_stride_1,
+            full_to_swa_mapping,
+            full_to_swa_mapping_stride_0,
+            batch_size,
+            max_seq_pages,
+            page_size,
+            seq_len_delta,
+            use_swa,
+            shift,
+            BLOCK_COLS=BLOCK_COLS,
+            num_warps=4,
+            num_stages=3,
+        )

python/sglang/kernels/ops/attention/pad.py

@@ -0,0 +1,390 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def pad_sequence_with_mask_kernel(
+    input_ptr,  # (total_tokens, hidden)
+    offsets_ptr,  # (B,)
+    lengths_ptr,  # (B,)
+    output_ptr,  # (B, max_len, hidden)
+    mask_ptr,  # (B, max_len)
+    max_len,
+    hidden_dim,
+    BLOCK_M: tl.constexpr,  # seq block
+    BLOCK_D: tl.constexpr,  # hidden block
+):
+    b = tl.program_id(0)  # batch index
+    m = tl.program_id(1)  # seq block index
+
+    offset = tl.load(offsets_ptr + b)
+    length = tl.load(lengths_ptr + b)
+
+    seq_ids = m * BLOCK_M + tl.arange(0, BLOCK_M)
+    hid_ids = tl.arange(0, BLOCK_D)
+
+    seq_mask = seq_ids < max_len
+    valid_token = seq_ids < length
+
+    # input index
+    in_token = offset + seq_ids
+    in_ptr = input_ptr + in_token[:, None] * hidden_dim + hid_ids[None, :]
+
+    # output index
+    out_ptr = (
+        output_ptr
+        + b * max_len * hidden_dim
+        + seq_ids[:, None] * hidden_dim
+        + hid_ids[None, :]
+    )
+
+    values = tl.load(
+        in_ptr,
+        mask=valid_token[:, None] & (hid_ids[None, :] < hidden_dim),
+        other=0.0,
+    )
+
+    tl.store(
+        out_ptr,
+        values,
+        mask=seq_mask[:, None] & (hid_ids[None, :] < hidden_dim),
+    )
+
+    # attention mask
+    if tl.program_id(2) == 0:
+        mask_out_ptr = mask_ptr + b * max_len + seq_ids
+        tl.store(mask_out_ptr, valid_token, mask=seq_mask)
+
+
+def pad_sequence_with_mask(
+    input_emb,  # (total_tokens, hidden)
+    offsets,  # (B,)
+    lengths,  # (B,)
+    max_len,
+):
+    B = offsets.shape[0]
+    hidden_dim = input_emb.shape[1]
+
+    output = torch.zeros(
+        (B, max_len, hidden_dim),
+        device=input_emb.device,
+        dtype=input_emb.dtype,
+    )
+    attn_mask = torch.empty(
+        (B * max_len),
+        device=input_emb.device,
+        dtype=torch.bool,
+    )
+
+    BLOCK_D = triton.next_power_of_2(hidden_dim)
+    BLOCK_M = triton.next_power_of_2(max_len)
+
+    grid = (
+        B,
+        triton.cdiv(max_len, BLOCK_M),
+        1,
+    )
+
+    pad_sequence_with_mask_kernel[grid](
+        input_emb,
+        offsets,
+        lengths,
+        output,
+        attn_mask,
+        max_len,
+        hidden_dim,
+        BLOCK_M=BLOCK_M,
+        BLOCK_D=BLOCK_D,
+    )
+
+    return B, output, attn_mask
+
+
+@triton.jit
+def pad_draft_extend_query_kernel(
+    q_ptr,  # Input query tensor [total_seq_len, num_heads, head_dim]
+    padded_q_ptr,  # Output padded query tensor [batch_size, max_seq_len, num_heads, head_dim]
+    seq_lens_q_ptr,  # Sequence lengths for each sequence [batch_size]
+    cumsum_ptr,  # Cumulative sum of sequence lengths [batch_size + 1]
+    batch_size,
+    max_seq_len,
+    num_heads,
+    head_dim,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Triton kernel for padding draft extended query tensor with parallelized head and dim processing."""
+    # Use 3D program IDs: (batch_seq, head_block, dim_block)
+    batch_seq_pid = tl.program_id(0)
+    head_pid = tl.program_id(1)
+    dim_pid = tl.program_id(2)
+
+    batch_id = batch_seq_pid // max_seq_len
+    seq_pos = batch_seq_pid % max_seq_len
+
+    if batch_id >= batch_size:
+        return
+
+    # Load sequence length for this batch
+    seq_len = tl.load(seq_lens_q_ptr + batch_id)
+
+    if seq_pos >= seq_len:
+        return
+
+    # Load cumulative sum to get start position in input tensor
+    input_start = tl.load(cumsum_ptr + batch_id)
+    input_pos = input_start + seq_pos
+
+    # Calculate head and dim block ranges
+    head_start = head_pid * BLOCK_SIZE
+    head_end = tl.minimum(head_start + BLOCK_SIZE, num_heads)
+    head_mask = tl.arange(0, BLOCK_SIZE) < (head_end - head_start)
+
+    dim_start = dim_pid * BLOCK_SIZE
+    dim_end = tl.minimum(dim_start + BLOCK_SIZE, head_dim)
+    dim_mask = tl.arange(0, BLOCK_SIZE) < (dim_end - dim_start)
+
+    # Calculate input offset
+    input_offset = (
+        input_pos * num_heads * head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Load data
+    data = tl.load(
+        q_ptr + input_offset,
+        mask=head_mask[:, None] & dim_mask[None, :],
+        other=0.0,
+    )
+
+    # Calculate output offset
+    output_offset = (
+        batch_id * max_seq_len * num_heads * head_dim
+        + seq_pos * num_heads * head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Store data
+    tl.store(
+        padded_q_ptr + output_offset,
+        data,
+        mask=head_mask[:, None] & dim_mask[None, :],
+    )
+
+
+def pad_draft_extend_query(
+    q: torch.Tensor,
+    padded_q: torch.Tensor,
+    seq_lens_q: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+) -> torch.Tensor:
+    """Pad draft extended query using Triton kernel."""
+    batch_size = cu_seqlens_q.shape[0] - 1
+    max_seq_len_q = padded_q.shape[1]
+    num_heads = padded_q.shape[2]
+    head_dim = padded_q.shape[3]
+
+    # Launch Triton kernel with 3D grid for parallelized head and dim processing
+    BLOCK_SIZE = 64
+    num_head_blocks = triton.cdiv(num_heads, BLOCK_SIZE)
+    num_dim_blocks = triton.cdiv(head_dim, BLOCK_SIZE)
+    grid = (batch_size * max_seq_len_q, num_head_blocks, num_dim_blocks)
+
+    pad_draft_extend_query_kernel[grid](
+        q_ptr=q,
+        padded_q_ptr=padded_q,
+        seq_lens_q_ptr=seq_lens_q,
+        cumsum_ptr=cu_seqlens_q,
+        batch_size=batch_size,
+        max_seq_len=max_seq_len_q,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return padded_q
+
+
+@triton.jit
+def unpad_draft_extend_output_kernel(
+    raw_out_ptr,  # Input raw output tensor (batch_size, token_per_batch, tp_q_head_num, v_head_dim)
+    output_ptr,  # Output tensor (-1, tp_q_head_num, v_head_dim)
+    num_accept_tokens_ptr,  # Accept lengths for each sequence [batch_size]
+    cumsum_ptr,  # Cumulative sum of accept lengths [batch_size + 1]
+    batch_size,
+    token_per_batch,
+    tp_q_head_num,
+    v_head_dim,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Triton kernel for unpadding draft extended output tensor with parallelized head and dim processing."""
+    batch_seq_pid = tl.program_id(0)
+    head_pid = tl.program_id(1)
+    dim_pid = tl.program_id(2)
+
+    batch_id = batch_seq_pid // token_per_batch
+    seq_pos = batch_seq_pid % token_per_batch
+
+    if batch_id >= batch_size:
+        return
+
+    # Load accept length for this batch
+    accept_len = tl.load(num_accept_tokens_ptr + batch_id)
+
+    if seq_pos >= accept_len:
+        return
+
+    # Load cumulative sum to get start position in output tensor
+    output_start = tl.load(cumsum_ptr + batch_id)
+    output_pos = output_start + seq_pos
+
+    # Calculate head and dim block ranges
+    head_start = head_pid * BLOCK_SIZE
+    head_end = tl.minimum(head_start + BLOCK_SIZE, tp_q_head_num)
+    head_mask = tl.arange(0, BLOCK_SIZE) < (head_end - head_start)
+
+    dim_start = dim_pid * BLOCK_SIZE
+    dim_end = tl.minimum(dim_start + BLOCK_SIZE, v_head_dim)
+    dim_mask = tl.arange(0, BLOCK_SIZE) < (dim_end - dim_start)
+
+    # Calculate input offset: (batch_id, seq_pos, head_id, dim_id)
+    input_offset = (
+        batch_id * token_per_batch * tp_q_head_num * v_head_dim
+        + seq_pos * tp_q_head_num * v_head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * v_head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Load data
+    data = tl.load(
+        raw_out_ptr + input_offset,
+        mask=head_mask[:, None] & dim_mask[None, :],
+        other=0.0,
+    )
+
+    output_offset = (
+        output_pos * tp_q_head_num * v_head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * v_head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Store data
+    tl.store(
+        output_ptr + output_offset,
+        data,
+        mask=head_mask[:, None] & dim_mask[None, :],
+    )
+
+
+def unpad_draft_extend_output(
+    raw_out: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    seq_lens_q: torch.Tensor,
+    sum_seq_lens_q: int,
+    unpad_output_buffer: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Unpad draft extended output using Triton kernel."""
+    # raw_out: (batch_size, token_per_batch, layer.tp_q_head_num, layer.v_head_dim)
+    batch_size = seq_lens_q.shape[0]
+    token_per_batch = raw_out.shape[1]  # max_seq_len
+    tp_q_head_num = raw_out.shape[2]  # num_heads
+    v_head_dim = raw_out.shape[3]  # head_dim
+    total_tokens = sum_seq_lens_q
+
+    # Check if we're in CUDA graph mode (buffers are pre-allocated)
+    if unpad_output_buffer is not None:
+        # Use pre-allocated buffer for CUDA graph compatibility
+        output = unpad_output_buffer[:total_tokens, :, :].to(dtype=raw_out.dtype)
+    else:
+        # Dynamic allocation for non-CUDA graph mode
+        output = torch.empty(
+            (total_tokens, tp_q_head_num, v_head_dim),
+            dtype=raw_out.dtype,
+            device=raw_out.device,
+        )
+
+    # Launch Triton kernel with 3D grid for parallelized head and dim processing
+    BLOCK_SIZE = 64
+    num_head_blocks = triton.cdiv(tp_q_head_num, BLOCK_SIZE)
+    num_dim_blocks = triton.cdiv(v_head_dim, BLOCK_SIZE)
+    grid = (batch_size * token_per_batch, num_head_blocks, num_dim_blocks)
+
+    unpad_draft_extend_output_kernel[grid](
+        raw_out_ptr=raw_out,
+        output_ptr=output,
+        num_accept_tokens_ptr=seq_lens_q,
+        cumsum_ptr=cu_seqlens_q,
+        batch_size=batch_size,
+        token_per_batch=token_per_batch,
+        tp_q_head_num=tp_q_head_num,
+        v_head_dim=v_head_dim,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return output[:total_tokens, :, :]
+
+
+@triton.jit
+def seqlens_expand_kernel(
+    extend_seq_lens_ptr,  # [N]
+    seq_lens_ptr,  # [N]
+    offsets_ptr,  # [N+1]
+    output_ptr,  # [sum(extend_seq_lens)]
+    N,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    if pid >= N:
+        return
+
+    qo_len = tl.load(extend_seq_lens_ptr + pid)
+    kv_len = tl.load(seq_lens_ptr + pid)
+
+    start = kv_len - qo_len + 1
+    out_offset = tl.load(offsets_ptr + pid)
+
+    offs = tl.arange(0, BLOCK)
+    mask = offs < qo_len
+
+    # Clamp to >= 0: rows with kv_len < qo_len (DP-padded / idle-companion
+    # rows whose kv is the CUDA-graph fill value) would otherwise produce
+    # negative lengths, which unsigned consumers (e.g. the top-k v2 kernel,
+    # which reads lengths as uint32) turn into ~4e9-token lengths and an
+    # illegal memory access.
+    values = tl.maximum(start + offs, 0)
+    tl.store(output_ptr + out_offset + offs, values, mask=mask)
+
+
+def seqlens_expand_triton(
+    extend_seq_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    total_len: int,
+    max_q_len: int,
+):
+    """
+    extend_seq_lens: [N], int32, CUDA
+    seq_lens:        [N], int32, CUDA
+    """
+    assert extend_seq_lens.is_cuda
+    assert seq_lens.is_cuda
+
+    N = extend_seq_lens.numel()
+
+    offsets = torch.zeros(N + 1, device=extend_seq_lens.device, dtype=torch.int32)
+    offsets[1:] = torch.cumsum(extend_seq_lens, dim=0)
+    output = torch.empty(total_len, device=extend_seq_lens.device, dtype=torch.int32)
+
+    BLOCK = triton.next_power_of_2(max_q_len)
+    grid = (N,)
+
+    seqlens_expand_kernel[grid](
+        extend_seq_lens,
+        seq_lens,
+        offsets,
+        output,
+        N,
+        BLOCK=BLOCK,
+    )
+
+    return output

python/sglang/kernels/ops/attention/position.py

@@ -0,0 +1,59 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def compute_position_triton(
+    extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_sum
+):
+    """Compute positions. It is a fused version of `compute_position_torch`."""
+    batch_size = extend_seq_lens.shape[0]
+    has_prefix = extend_prefix_lens.shape[0] == batch_size
+
+    positions = torch.empty(
+        extend_seq_lens_sum, dtype=torch.int64, device=extend_seq_lens.device
+    )
+    extend_start_loc = torch.empty(
+        batch_size, dtype=torch.int32, device=extend_seq_lens.device
+    )
+
+    # Launch kernel
+    compute_position_kernel[(batch_size,)](
+        positions,
+        extend_start_loc,
+        extend_prefix_lens,
+        extend_seq_lens,
+        has_prefix,
+    )
+
+    return positions, extend_start_loc
+
+
+@triton.jit
+def compute_position_kernel(
+    positions,
+    extend_start_loc,
+    extend_prefix_lens,
+    extend_seq_lens,
+    has_prefix: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(0).to(tl.int64)
+
+    prefix_len = tl.load(extend_prefix_lens + pid) if has_prefix else 0
+    seq_len = tl.load(extend_seq_lens + pid)
+
+    # NOTE: This can be slow for large bs
+    cumsum_start = tl.cast(0, tl.int64)
+    for i in range(pid):
+        cumsum_start += tl.load(extend_seq_lens + i)
+
+    num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        tl.store(
+            positions + cumsum_start + offset,
+            prefix_len + offset,
+            mask=offset < seq_len,
+        )
+    tl.store(extend_start_loc + pid, cumsum_start)

python/sglang/kernels/ops/attention/prefill_attention.py

@@ -0,0 +1,219 @@
+# 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.
+# ==============================================================================
+"""
+Memory-efficient attention for prefill.
+It supporst page size = 1.
+"""
+
+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cuda, is_hip
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+if _is_cuda or _is_hip:
+    CUDA_CAPABILITY = torch.cuda.get_device_capability()
+
+
+@triton.jit
+def _fwd_kernel(
+    Q,
+    K,
+    V,
+    sm_scale,
+    B_Start_Loc,
+    B_Seqlen,
+    Out,
+    stride_qbs,
+    stride_qh,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    stride_obs,
+    stride_oh,
+    kv_group_num: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    Lk: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    start_m = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+
+    block_start_loc = BLOCK_M * start_m
+
+    # initialize offsets
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_q = (
+        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
+    off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :]
+
+    mask_d = offs_d < Lk
+
+    q = tl.load(
+        Q + off_q,
+        mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]),
+        other=0.0,
+    )
+
+    k_ptrs = K + off_k
+    v_ptrs = V + off_v
+
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+    block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
+
+    end_n = (
+        cur_batch_seq_len
+        if not IS_CAUSAL
+        else tl.minimum((start_m + 1) * BLOCK_M, cur_batch_seq_len)
+    )
+    for start_n in range(0, block_mask * end_n, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        # -- compute qk ----
+        k = tl.load(
+            k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
+            mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (mask_d[:, None]),
+            other=0.0,
+        )
+        # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0)
+
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk += tl.dot(q, k)
+        qk *= sm_scale
+
+        if IS_CAUSAL:
+            qk += tl.where(
+                (start_n + offs_n[None, :] < cur_batch_seq_len)
+                & (offs_m[:, None] >= (start_n + offs_n[None, :])),
+                0,
+                float("-inf"),
+            )
+        else:
+            qk += tl.where(
+                (start_n + offs_n[None, :]) < cur_batch_seq_len, 0, float("-inf")
+            )
+
+        # -- compute m_ij, p, l_ij
+        m_ij = tl.max(qk, 1)
+        p = tl.exp(qk - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        m_i_new = tl.maximum(m_i, m_ij)
+        alpha = tl.exp(m_i - m_i_new)
+        beta = tl.exp(m_ij - m_i_new)
+        l_i_new = alpha * l_i + beta * l_ij
+        # -- update output accumulator --
+        # scale p
+        p_scale = beta / l_i_new
+        p = p * p_scale[:, None]
+        # scale acc
+        acc_scale = l_i / l_i_new * alpha
+        acc = acc * acc_scale[:, None]
+        # update acc
+        v = tl.load(
+            v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
+            mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]),
+            other=0.0,
+        )
+
+        p = p.to(v.dtype)
+        acc += tl.dot(p, v)
+        # update m_i and l_i
+        l_i = l_i_new
+        m_i = m_i_new
+    # initialize pointers to output
+    off_o = (
+        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
+        + cur_head * stride_oh
+        + offs_d[None, :]
+    )
+    out_ptrs = Out + off_o
+    tl.store(
+        out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :])
+    )
+
+
+def context_attention_fwd(
+    q, k, v, o, b_start_loc, b_seq_len, max_input_len, is_causal=True, sm_scale=None
+):
+    """
+    q, k, v: [b * s, head, head_dim]
+    b_start_loc: [b]
+    b_seq_len: [b]
+    out: [b * s, head, head_dim]
+    sm_scale: softmax scale, defaults to 1/sqrt(head_dim)
+    """
+    if (_is_cuda or _is_hip) and CUDA_CAPABILITY[0] > 8:
+        BLOCK = 128
+    else:
+        BLOCK = 64
+
+    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+
+    if sm_scale is None:
+        sm_scale = 1.0 / (Lq**0.5)
+    batch, head = b_seq_len.shape[0], q.shape[1]
+    kv_group_num = q.shape[1] // k.shape[1]
+
+    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
+    num_warps = 4 if Lk <= 64 else 8
+
+    _fwd_kernel[grid](
+        q,
+        k,
+        v,
+        sm_scale,
+        b_start_loc,
+        b_seq_len,
+        o,
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        o.stride(0),
+        o.stride(1),
+        kv_group_num=kv_group_num,
+        BLOCK_M=BLOCK,
+        BLOCK_DMODEL=triton.next_power_of_2(Lk),
+        BLOCK_N=BLOCK,
+        IS_CAUSAL=is_causal,
+        num_warps=num_warps,
+        num_stages=1,
+        Lk=Lk,
+    )

python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py

@@ -0,0 +1,439 @@
+# 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.
+# ==============================================================================
+"""
+Memory-efficient attention for decoding.
+It supports page size = 1.
+"""
+
+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
+
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.attention.decode_attention import (
+    _decode_softmax_reducev_fwd,
+)
+
+
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+
+_is_hip = is_hip()
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.jit
+def _fwd_grouped_kernel_stage1_rope(
+    Q,  # Holds [Q_NOPE; Q_PE], b x h x (d+r)
+    K_Buffer,  # Holds [KV; K_PE], b*s x (c+r)
+    V_buffer,  # Holds [KV], b*s x (c)
+    cos_sin_cache,  # max_seq_len x (rotary_dim * 2)
+    positions,  # sequence positions
+    sm_scale,
+    kv_indptr,
+    kv_indices,
+    Att_Out,  # b x h x NUM_KV_SPLITS x (kv_lora_rank + 1)
+    k_pe_t_out,
+    stride_qb,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_vbs,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_kpe_tokens_out_b,
+    stride_cos_sin_cache_s,
+    stride_positions_b,
+    rotary_dim: tl.constexpr,
+    kv_lora_rank: tl.constexpr,
+    qk_rope_head_dim: tl.constexpr,
+    kv_group_num: tl.constexpr,
+    q_head_num: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    NUM_KV_SPLITS: tl.constexpr,
+    logit_cap: tl.constexpr,
+    USE_ROPE: tl.constexpr,
+    IS_NEOX_STYLE: tl.constexpr,
+):
+
+    cur_batch = tl.program_id(0)
+    cur_head_id = tl.program_id(1)
+    split_kv_id = tl.program_id(2)
+
+    if BLOCK_H < kv_group_num:
+        VALID_BLOCK_H: tl.constexpr = BLOCK_H
+    else:
+        VALID_BLOCK_H: tl.constexpr = kv_group_num
+    cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
+    mask_h = mask_h & (cur_head < q_head_num)
+
+    offs_c = tl.arange(0, BLOCK_C)
+    offs_qk_r = tl.arange(kv_lora_rank, kv_lora_rank + BLOCK_R)  # to get the k_pe
+
+    off_q_pe = (
+        cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_qk_r[None, :]
+    )
+    offs_q = cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_c[None, :]
+
+    mask_c = offs_c < kv_lora_rank
+    mask_qk_r = offs_qk_r < (kv_lora_rank + qk_rope_head_dim)
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+
+    q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_c[None, :]), other=0.0)
+    q_pe = tl.load(
+        Q + off_q_pe, mask=(mask_h[:, None]) & (mask_qk_r[None, :]), other=0.0
+    )
+
+    kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    # apply rotary embedding for q_pe, and k_pe (last token per batch of K_PE)
+    LAST_SPLIT = split_kv_end == cur_batch_seq_len
+    k_pe_last_token = tl.zeros([BLOCK_R], dtype=q.dtype)
+
+    if USE_ROPE:
+        if IS_NEOX_STYLE:
+            # [BLOCK_ROTARY // 2, BLOCK_ROTARY // 2 + 1, BLOCK_ROTARY // 2 + 2, ..., 0, 1, 2, ..., BLOCK_ROTARY // 2 - 1, pass:]
+            offs_qk_rot_r = kv_lora_rank + (
+                (tl.arange(0, BLOCK_R) + (rotary_dim // 2)) % rotary_dim
+            )
+            # Which elements to flip
+            mask_rotate = tl.arange(0, BLOCK_R) < (rotary_dim // 2)
+            # [0 , 1, 2, ..., rotary_dim // 2 - 1, 0 , 1, 2, ..., rotary_dim // 2 - 1]
+            offs_rotary = tl.arange(0, BLOCK_R) % (rotary_dim // 2)
+        else:
+            # [1, 0, 3, 2, 5, 4, ..., BLOCK_R, BLOCK_R - 1]
+            offs_qk_rot_r = (
+                kv_lora_rank
+                + (((tl.arange(0, BLOCK_R) + 1) % 2) * 2)
+                - 1
+                + tl.arange(0, BLOCK_R)
+            )
+            mask_rotate = tl.arange(0, BLOCK_R) % 2 < 1
+            # [0, 0, 1, 1, ..., rotary_dim // 2 - 1, rotary_dim // 2 - 1]
+            offs_rotary = tl.arange(0, BLOCK_R) // 2
+
+        if qk_rope_head_dim > rotary_dim:
+            offs_qk_rot_r = tl.where(
+                tl.arange(0, BLOCK_R) < rotary_dim, offs_qk_rot_r, tl.arange(0, BLOCK_R)
+            )
+            offs_rotary = tl.where(
+                tl.arange(0, BLOCK_R) < rotary_dim, offs_rotary, tl.arange(0, BLOCK_R)
+            )
+
+        mask_rotary = tl.arange(0, BLOCK_R) < rotary_dim
+
+        pos = tl.load(positions + cur_batch * stride_positions_b)
+        cos = tl.load(
+            cos_sin_cache + pos * stride_cos_sin_cache_s + offs_rotary,
+            mask=mask_rotary,
+            other=1.0,
+        )
+        sin = tl.load(
+            cos_sin_cache
+            + pos * stride_cos_sin_cache_s
+            + offs_rotary
+            + rotary_dim // 2,
+            mask_rotary,
+            other=0.0,
+        )
+
+        off_q_pe_rot = (
+            cur_batch * stride_qb
+            + cur_head[:, None] * stride_qh
+            + offs_qk_rot_r[None, :]
+        )
+        mask_qk_rot_r = offs_qk_rot_r < (kv_lora_rank + qk_rope_head_dim)
+
+        # 0, 2, 4,.... 1, 3, 5...
+        q_pe_rot = tl.load(
+            Q + off_q_pe_rot,
+            mask=(mask_h[:, None]) & (mask_qk_rot_r[None, :]),
+            other=0.0,
+        )
+        q_pe_rot = tl.where(mask_rotate[None, :], -q_pe_rot, q_pe_rot)
+
+        q_pe = q_pe * cos + q_pe_rot * sin
+
+        # we only apply to the last token in the K_PE
+        if LAST_SPLIT:
+            # debug assert
+            if (cur_batch == 0 and cur_head == 0) and split_kv_id < NUM_KV_SPLITS - 1:
+                tl.device_assert(False, "Only last split should compute k_pe")
+
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + cur_batch_seq_len - 1
+            )
+            offs_buf_k_pe_last_token = kv_loc * stride_buf_kbs + offs_qk_r
+            offs_buf_k_pe_rot_last_token = kv_loc * stride_buf_kbs + offs_qk_rot_r
+            k_pe_last_token = tl.load(K_Buffer + offs_buf_k_pe_last_token)
+
+            k_pe_rot_last_token = tl.load(K_Buffer + offs_buf_k_pe_rot_last_token)
+            k_pe_rot_last_token = tl.where(
+                mask_rotate, -k_pe_rot_last_token, k_pe_rot_last_token
+            )
+
+            k_pe_last_token = k_pe_last_token * cos + k_pe_rot_last_token * sin
+
+    e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, BLOCK_C], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+
+            offs_buf_kv = kv_loc[None, :] * stride_buf_kbs + offs_c[:, None]
+            offs_buf_k_pe = kv_loc[None, :] * stride_buf_kbs + offs_qk_r[:, None]
+
+            k_pe = tl.load(
+                K_Buffer + offs_buf_k_pe,
+                mask=(offs_n[None, :] < split_kv_end) & (mask_qk_r[:, None]),
+                other=0.0,
+            )  # positional embedding part of keys
+
+            if (USE_ROPE and LAST_SPLIT) and start_n >= cur_batch_seq_len - BLOCK_N:
+                k_pe = tl.where(
+                    offs_n[None, :] != (split_kv_end - 1),
+                    k_pe,
+                    k_pe_last_token[:, None],
+                )
+
+            # (16, 64) x (64, 32)
+            # dot product of rope parts
+            qk = tl.dot(q_pe, k_pe.to(q_pe.dtype))
+
+            kv = tl.load(
+                K_Buffer + offs_buf_kv,
+                mask=(offs_n[None, :] < split_kv_end) & (mask_c[:, None]),
+                other=0.0,
+            )  # the shared latent tensor for keys and values
+
+            # (16, 512) x (512, 32)
+            # dot product of nope parts
+            qk += tl.dot(q, kv)
+
+            qk *= sm_scale
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            qk = tl.where(
+                mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
+            )
+
+            offs_buf_v = kv_loc[:, None] * stride_buf_vbs + offs_c[None, :]
+            v = tl.load(
+                V_buffer + offs_buf_v,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_c[None, :]),
+                other=0.0,
+            )
+
+            n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max[:, None])
+            acc *= re_scale[:, None]
+            # (16, 32) x (32, 512)
+            acc += tl.dot(p.to(v.dtype), v)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 1)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head[:, None] * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_c[None, :]
+        )
+
+        if USE_ROPE:
+            if LAST_SPLIT:
+                k_pe_last_token_ptrs = (
+                    k_pe_t_out
+                    + cur_batch * stride_kpe_tokens_out_b
+                    + tl.arange(0, BLOCK_R)
+                )
+                tl.store(k_pe_last_token_ptrs, k_pe_last_token, mask=mask_qk_r)
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum[:, None],
+            mask=(mask_h[:, None]) & (mask_c[None, :]),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + kv_lora_rank
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+            mask=mask_h,
+        )
+
+
+# TODO rope offset
+def _decode_grouped_att_m_fwd_rope(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    k_pe_tokens_out,
+    kv_lora_rank,  # c
+    cos_sin_cache,
+    positions,
+    rotary_dim,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    sm_scale,
+    logit_cap,
+    use_rope,
+    is_neox_style=True,
+):
+    if use_rope:
+        assert (
+            k_pe_tokens_out is not None
+        ), "We must output the k_pe tokens with rope applied if rope fusion enabled."
+
+    BLOCK = 32
+
+    # # [TODO] work around shmem limit on MI3xx
+    # if _is_hip and kv_lora_rank >= 576:
+    #     BLOCK = 16
+
+    qk_rope_head_dim = k_buffer.shape[-1] - kv_lora_rank
+    batch, head_num = kv_indptr.shape[0] - 1, q.shape[1]
+    kv_group_num = q.shape[1] // k_buffer.shape[1]
+
+    BLOCK_C = triton.next_power_of_2(kv_lora_rank)
+    BLOCK_R = triton.next_power_of_2(qk_rope_head_dim)
+
+    BLOCK_H = 16
+    NUM_KV_SPLITS = num_kv_splits
+    grid = (
+        batch,
+        triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
+        NUM_KV_SPLITS,
+    )
+
+    extra_kargs = {}
+    num_stages = 2
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
+        num_stages = 1
+
+    _fwd_grouped_kernel_stage1_rope[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        cos_sin_cache,
+        positions,
+        sm_scale,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        k_pe_tokens_out,
+        q.stride(0),
+        q.stride(1),
+        k_buffer.stride(0),
+        v_buffer.stride(0),
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        k_pe_tokens_out.stride(0) if use_rope else 0,
+        cos_sin_cache.stride(0) if use_rope else 0,
+        positions.stride(0) if use_rope else 0,
+        rotary_dim,
+        kv_lora_rank,
+        qk_rope_head_dim,
+        kv_group_num=kv_group_num,
+        q_head_num=head_num,
+        BLOCK_C=BLOCK_C,
+        BLOCK_R=BLOCK_R,
+        BLOCK_N=BLOCK,
+        BLOCK_H=BLOCK_H,
+        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        logit_cap=logit_cap,
+        USE_ROPE=use_rope,
+        IS_NEOX_STYLE=is_neox_style,
+        num_warps=4,
+        num_stages=num_stages,
+        **extra_kargs,
+    )
+
+
+def decode_attention_fwd_grouped_rope(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    k_pe_tokens,
+    kv_lora_rank,
+    rotary_dim,
+    cos_sin_cache,
+    positions,
+    attn_logits,
+    num_kv_splits,
+    sm_scale,
+    logit_cap=0.0,
+    use_rope=False,
+    is_neox_style=False,
+):
+    _decode_grouped_att_m_fwd_rope(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        k_pe_tokens,
+        kv_lora_rank,
+        cos_sin_cache,
+        positions,
+        rotary_dim,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        sm_scale,
+        logit_cap,
+        use_rope,
+        is_neox_style,
+    )
+    _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, kv_indptr, num_kv_splits)

python/sglang/kernels/ops/attention/verify_splitkv.py

@@ -0,0 +1,803 @@
+"""Split-KV (flash-decode) attention for EAGLE speculative *verify*.
+
+Only valid when speculative ``topk == 1`` (the EAGLE tree reduces to a pure
+causal chain); the caller gates on that. ``topk > 1`` trees fall back to
+``extend_attention_fwd``.
+
+On the Triton backend, EAGLE target-verify runs through the prefill
+``extend_attention_fwd``, which loops the (long) prefix KV serially per
+(sequence, head). With only a few draft-token queries, that leaves the GPU
+memory system far under-utilized at long context. This kernel instead splits
+the prefix KV across parallel programs (flash-decode style) and combines the
+partials with a log-sum-exp merge, then handles the small causal draft-draft
+block -- recovering memory bandwidth on the verify path.
+
+Two Triton kernels:
+  * ``_verify_prefix_stage1``: split-KV over the shared prefix. Applies the fp8
+    dequant multipliers ``k_scale`` (on the QK score) and ``v_scale`` (on the
+    prefix output), matching ``extend_attention_fwd``'s ``_fwd_kernel``
+    (qk *= sm_scale * k_scale; acc += dot(p, v) * v_scale on the prefix loop;
+    NO scaling on the draft-draft loop, whose K/V are the fresh bf16 draft
+    tensors, not the fp8 pool). fp8 K/V buffers are handled by casting q to the
+    buffer dtype before the dot (mirrors ``q.to(k.dtype)`` in the baseline).
+  * ``_verify_combine_stage2``: combines the prefix splits (LSE merge) with the
+    small causal draft-draft block and writes the output.
+
+``verify_splitkv_fwd(...)`` takes the SAME positional args as
+``extend_attention_fwd``; it runs the split-KV path when it can serve the case
+bit-equivalently and returns True, otherwise returns False (doing nothing) so
+the caller falls back to ``extend_attention_fwd``. Supported case: causal
+(topk=1) verify with a constant per-sequence extend length, no sinks /
+sliding-window / logit-cap / xai-temperature. Correctness is never violated.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_hip
+
+_MIN_BLOCK_KV = 32
+
+# AMD/CDNA-only Triton launch hints (waves_per_eu, matrix_instr_nonkdim); NVIDIA's
+# Triton rejects these kwargs, so only pass them on ROCm. In production this kernel
+# is dispatched only on AMD (see TritonAttnBackend); keeping it NV-safe lets the
+# numerics test run on the CUDA CI lane.
+_IS_HIP = is_hip()
+_AMD_LAUNCH_KWARGS = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16} if _IS_HIP else {}
+
+# Block-size config keyed on head_dim. The (BLOCK_N, num_warps) tile that best
+# hides latency depends on head_dim: at head_dim=256 (Qwen3 family) a narrower
+# BLOCK_N with more warps wins, since the 256-wide QK/PV tiles are register
+# heavy. head_dim=256 is the value validated on MI350X; other head dims use a
+# conservative default. Block size affects PERFORMANCE only, never correctness
+# (any valid block size produces the same result).
+DEFAULT_N_SPLITS = 8
+DEFAULT_BLOCK_N = 32
+DEFAULT_NUM_WARPS = 4
+_BLOCK_CONFIG = {
+    # head_dim: (BLOCK_N, num_warps)
+    256: (32, 4),
+}
+
+
+def block_config(head_dim):
+    """Return (BLOCK_N, num_warps) for a head_dim; default for untuned dims."""
+    return _BLOCK_CONFIG.get(head_dim, (DEFAULT_BLOCK_N, DEFAULT_NUM_WARPS))
+
+
+# ---------------------------------------------------------------------------
+# Adaptive N_SPLITS.
+# ---------------------------------------------------------------------------
+# The prefix split-KV stage launches a (bs, h_q, N_SPLITS) grid; each (b,h,s)
+# program handles kv_len_per_split = cdiv(cdiv(seqlen, N_SPLITS), MIN)*MIN keys.
+# A fixed N_SPLITS=16 over-splits short/mid contexts (each split does too little
+# work -> launch + reduction overhead dominates) and under-splits very long ones
+# (too few parallel waves to saturate the device, raising tail latency on the
+# slow split). Mirror the decode kernel's intent (decode_attention.py
+# get_num_kv_splits): pick the split count per-dispatch from the representative
+# sequence length, growing gradually with seqlen and capped at MAX.
+#
+# CRITICAL: this must be computed from STATIC shapes only (no .item()/.cpu()
+# sync), because the verify/draft-extend step runs inside a captured HIP graph
+# where a device->host copy raises hipErrorStreamCaptureUnsupported. We use the
+# average prefix length = kv_indices.shape[0] / bs, which is a pure python int
+# from tensor shapes -- no device read. N_SPLITS is then a power of two so the
+# stage2 reduction tile (tl.arange(0, N_SPLITS)) stays cheap.
+#
+# Split-count bounds (internal constants). MAX=16 is the MI350X cap: 32
+# oversubscribes the device and regresses, per tuning.
+ADAPTIVE_SPLITS = True
+MAX_N_SPLITS = 16
+MIN_N_SPLITS = 4
+
+
+def choose_n_splits(avg_seqlen):
+    """Pick N_SPLITS (power of two, in [MIN_N_SPLITS, MAX_N_SPLITS]) from the
+    average prefix length. Tuned by the real-shape sweep (head_dim=256, BS*H_Q
+    =128 base programs on ~132 CUs):
+
+        ctx  <  4k -> 4   (short: extra splits add launch/reduction overhead)
+        4k <= ctx < 8k -> 8   (sweet spot: best across 1k-16k in the sweep)
+        ctx >= 8k      -> 16  (long: a few more splits help latency-bound tail)
+
+    Never 32 (4096 grid blocks oversubscribes the device and regresses, per the
+    sweep). Computed from a static shape (avg prefix = kv_indices.shape[0]/bs),
+    so it is HIP-graph-capture safe (no device->host sync)."""
+    if not ADAPTIVE_SPLITS:
+        return DEFAULT_N_SPLITS
+    s = int(avg_seqlen)
+    if s < 4096:
+        n = 4
+    elif s < 8192:
+        n = 8
+    else:
+        n = 16
+    if n < MIN_N_SPLITS:
+        n = MIN_N_SPLITS
+    if n > MAX_N_SPLITS:
+        n = MAX_N_SPLITS
+    return n
+
+
+@triton.jit
+def _verify_prefix_stage1(
+    Q,  # [extend_tokens, H_Q, D]
+    K_Buffer,  # [pool_tokens, H_KV, D]
+    V_Buffer,  # [pool_tokens, H_KV, Dv]
+    sm_scale,
+    k_scale,  # fp8 dequant multiplier for prefix K (1.0 if bf16)
+    v_scale,  # fp8 dequant multiplier for prefix V (1.0 if bf16)
+    qo_indptr,  # [BS+1] int32  -> rows of Q (draft queries)
+    kv_indptr,  # [BS+1] int32  -> rows of kv_indices (prefix)
+    kv_indices,  # [sum prefix] int64
+    Att_Out,  # [BS, H_Q, N_SPLITS, L_EXT, Dv]  fp32
+    Att_Lse,  # [BS, H_Q, N_SPLITS, L_EXT]      fp32
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    stride_ob,
+    stride_oh,
+    stride_os,
+    stride_ol,
+    stride_lb,
+    stride_lh,
+    stride_ls,
+    kv_group_num: tl.constexpr,
+    N_SPLITS: tl.constexpr,
+    L_EXT: tl.constexpr,  # padded power-of-2 row tile (>= real l_ext)
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    split_kv_id = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    offs_l = tl.arange(0, L_EXT)
+
+    # real number of draft query tokens for this seq
+    cur_q_start = tl.load(qo_indptr + cur_batch)
+    l_ext = tl.load(qo_indptr + cur_batch + 1) - cur_q_start
+    mask_l = offs_l < l_ext
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+
+    # split sizing identical to the decode kernel
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, N_SPLITS), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = tl.zeros([L_EXT], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([L_EXT], dtype=tl.float32)
+    acc = tl.zeros([L_EXT, BLOCK_DV], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        # q tile: [L_EXT, D]
+        offs_q = (
+            (cur_q_start + offs_l)[:, None] * stride_qbs
+            + cur_head * stride_qh
+            + offs_d[None, :]
+        )
+        q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0)
+        q_k = q.to(K_Buffer.dtype.element_ty)
+
+        base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None]
+        base_offs_v = cur_kv_head * stride_buf_vh + offs_dv[None, :]
+
+        for start_n in tl.range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            n_mask = offs_n < split_kv_end
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=n_mask,
+                other=0,
+            )
+            # K block: [D, BLOCK_N]
+            offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
+            k = tl.load(K_Buffer + offs_buf_k, mask=n_mask[None, :], other=0.0)
+            qk = tl.dot(q_k, k)  # [L_EXT, BLOCK_N]
+            qk *= sm_scale * k_scale  # fp8 dequant of prefix K (k_scale==1 if bf16)
+            # NO causal mask: full prefix is visible to all draft tokens.
+            qk = tl.where(n_mask[None, :], qk, float("-inf"))
+
+            # V block: [BLOCK_N, Dv]
+            offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
+            v = tl.load(V_Buffer + offs_buf_v, mask=n_mask[:, None], other=0.0)
+
+            n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max[:, None])
+            acc *= re_scale[:, None]
+            acc += tl.dot(p.to(v.dtype), v)
+            e_sum = e_sum * re_scale + tl.sum(p, 1)
+            e_max = n_e_max
+
+        # fp8 dequant of prefix V: scale the accumulated (pre-normalised) output.
+        acc *= v_scale
+
+        offs_o = (
+            cur_batch * stride_ob
+            + cur_head * stride_oh
+            + split_kv_id * stride_os
+            + offs_l[:, None] * stride_ol
+            + offs_dv[None, :]
+        )
+        tl.store(Att_Out + offs_o, acc / e_sum[:, None], mask=mask_l[:, None])
+
+        offs_lse = (
+            cur_batch * stride_lb
+            + cur_head * stride_lh
+            + split_kv_id * stride_ls
+            + offs_l
+        )
+        tl.store(Att_Lse + offs_lse, e_max + tl.log(e_sum), mask=mask_l)
+    else:
+        # split did not run: write a sentinel lse so stage2 can ignore it.
+        offs_lse = (
+            cur_batch * stride_lb
+            + cur_head * stride_lh
+            + split_kv_id * stride_ls
+            + offs_l
+        )
+        tl.store(
+            Att_Lse + offs_lse,
+            tl.zeros([L_EXT], tl.float32) - float("inf"),
+            mask=mask_l,
+        )
+
+
+@triton.jit
+def _verify_combine_stage2(
+    Att_Out,  # [BS, H_Q, N_SPLITS, L_EXT, Dv]  fp32
+    Att_Lse,  # [BS, H_Q, N_SPLITS, L_EXT]      fp32
+    Q,  # [extend_tokens, H_Q, D]   (draft queries)
+    K_Extend,  # [extend_tokens, H_KV, D]
+    V_Extend,  # [extend_tokens, H_KV, Dv]
+    O_Out,  # [extend_tokens, H_Q, Dv]  (final, written)
+    sm_scale,
+    qo_indptr,  # [BS+1] int32
+    stride_ob,
+    stride_oh,
+    stride_os,
+    stride_ol,
+    stride_lb,
+    stride_lh,
+    stride_ls,
+    stride_qbs,
+    stride_qh,
+    stride_kebs,
+    stride_keh,
+    stride_vebs,
+    stride_veh,
+    stride_oobs,
+    stride_ooh,
+    kv_group_num: tl.constexpr,
+    N_SPLITS: tl.constexpr,
+    L_EXT: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    offs_l = tl.arange(0, L_EXT)
+    offs_s = tl.arange(0, N_SPLITS)
+
+    cur_q_start = tl.load(qo_indptr + cur_batch)
+    l_ext = tl.load(qo_indptr + cur_batch + 1) - cur_q_start
+    mask_l = offs_l < l_ext
+
+    # ---- (a) combine prefix splits (logsumexp) ----------------------------
+    # lse: [N_SPLITS, L_EXT]
+    offs_lse = (
+        cur_batch * stride_lb
+        + cur_head * stride_lh
+        + offs_s[:, None] * stride_ls
+        + offs_l[None, :]
+    )
+    lse = tl.load(offs_lse + Att_Lse)  # [N_SPLITS, L_EXT]
+    m_p = tl.max(lse, 0)  # [L_EXT]
+    w = tl.exp(lse - m_p[None, :])  # [N_SPLITS, L_EXT]; -inf->0
+    denom_p = tl.sum(w, 0)  # [L_EXT]
+
+    # weighted-sum of partial outputs: o_prefix[L_EXT, Dv]
+    # Att_Out[b,h,s,l,dv]
+    offs_ao = (
+        cur_batch * stride_ob
+        + cur_head * stride_oh
+        + offs_s[:, None, None] * stride_os
+        + offs_l[None, :, None] * stride_ol
+        + offs_dv[None, None, :]
+    )
+    ao = tl.load(offs_ao + Att_Out)  # [N_SPLITS, L_EXT, Dv]
+    o_prefix = tl.sum(ao * w[:, :, None], 0)  # [L_EXT, Dv]
+    o_prefix = o_prefix / denom_p[:, None]
+    lse_prefix = m_p + tl.log(denom_p)  # [L_EXT]
+
+    # ---- (b) draft-draft causal attention (L_EXT x L_EXT) -----------------
+    # load draft queries [L_EXT, D], draft K/V [L_EXT, D]/[L_EXT, Dv]
+    offs_q = (
+        (cur_q_start + offs_l)[:, None] * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0).to(tl.float32)
+
+    offs_ke = (
+        (cur_q_start + offs_l)[:, None] * stride_kebs
+        + cur_kv_head * stride_keh
+        + offs_d[None, :]
+    )
+    ke = tl.load(K_Extend + offs_ke, mask=mask_l[:, None], other=0.0).to(tl.float32)
+    offs_ve = (
+        (cur_q_start + offs_l)[:, None] * stride_vebs
+        + cur_kv_head * stride_veh
+        + offs_dv[None, :]
+    )
+    ve = tl.load(V_Extend + offs_ve, mask=mask_l[:, None], other=0.0).to(tl.float32)
+
+    # scores[i,j] = q_i . k_j  (i query, j key)  -> [L_EXT, L_EXT]
+    qk = tl.sum(q[:, None, :] * ke[None, :, :], 2) * sm_scale
+    # causal among drafts: query i sees key j iff j <= i, and both valid
+    causal = (offs_l[None, :] <= offs_l[:, None]) & mask_l[None, :] & mask_l[:, None]
+    qk = tl.where(causal, qk, float("-inf"))
+    m_d = tl.max(qk, 1)  # [L_EXT]
+    pd = tl.exp(qk - m_d[:, None])  # [L_EXT, L_EXT]
+    denom_d = tl.sum(pd, 1)  # [L_EXT]
+    o_draft = tl.sum(pd[:, :, None] * ve[None, :, :], 1)  # [L_EXT, Dv]
+    o_draft = o_draft / denom_d[:, None]
+    lse_draft = m_d + tl.log(denom_d)  # [L_EXT]
+
+    # ---- (c) final LSE merge (prefix vs draft) ----------------------------
+    m = tl.maximum(lse_prefix, lse_draft)
+    wp = tl.exp(lse_prefix - m)
+    wd = tl.exp(lse_draft - m)
+    o = (o_prefix * wp[:, None] + o_draft * wd[:, None]) / (wp + wd)[:, None]
+
+    offs_oo = (
+        (cur_q_start + offs_l)[:, None] * stride_oobs
+        + cur_head * stride_ooh
+        + offs_dv[None, :]
+    )
+    tl.store(O_Out + offs_oo, o.to(O_Out.dtype.element_ty), mask=mask_l[:, None])
+
+
+class VerifySplitKV:
+    """Pre-allocates scratch buffers for a problem shape and runs the split-KV
+    verify attention end to end (two Triton launches: prefix split-KV + fused
+    combine/draft/merge). Buffers are sized by ``max_bs`` (constant for the
+    server lifetime) and reused for every batch size <= max_bs, so their
+    addresses stay fixed (CUDA/HIP-graph safe) and GPU memory does not grow per
+    batch size. The kernel grid uses the actual per-call bs (<= max_bs)."""
+
+    def __init__(
+        self,
+        max_bs,
+        h_q,
+        h_kv,
+        head_dim,
+        v_head_dim,
+        l_ext,
+        device="cuda",
+        n_splits=DEFAULT_N_SPLITS,
+        block_n=DEFAULT_BLOCK_N,
+        num_warps=DEFAULT_NUM_WARPS,
+    ):
+        self.h_q = h_q
+        self.h_kv = h_kv
+        self.group = h_q // h_kv
+        self.head_dim = head_dim
+        self.v_head_dim = v_head_dim
+        self.l_ext = l_ext  # real draft tokens per seq (fixed == 4)
+        self.l_pad = triton.next_power_of_2(l_ext)
+        self.device = device
+        self.n_splits = n_splits
+        self.block_n = block_n
+        self.num_warps = num_warps
+        self._alloc(max_bs)
+
+    def _alloc(self, max_bs):
+        # prefix split partials (fp32), sized for the maximum batch size.
+        self.max_bs = max_bs
+        self.att_out = torch.empty(
+            (max_bs, self.h_q, self.n_splits, self.l_pad, self.v_head_dim),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        self.att_lse = torch.empty(
+            (max_bs, self.h_q, self.n_splits, self.l_pad),
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+    def grow_buffers(self, max_bs):
+        if max_bs > self.max_bs:
+            self._alloc(max_bs)
+
+    def _run_prefix_kernel(
+        self,
+        bs,
+        q_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        sm_scale,
+        k_scale,
+        v_scale,
+    ):
+        grid = (bs, self.h_q, self.n_splits)
+        _verify_prefix_stage1[grid](
+            q_extend,
+            k_buffer,
+            v_buffer,
+            sm_scale,
+            k_scale,
+            v_scale,
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            self.att_out,
+            self.att_lse,
+            q_extend.stride(0),
+            q_extend.stride(1),
+            k_buffer.stride(0),
+            k_buffer.stride(1),
+            v_buffer.stride(0),
+            v_buffer.stride(1),
+            self.att_out.stride(0),
+            self.att_out.stride(1),
+            self.att_out.stride(2),
+            self.att_out.stride(3),
+            self.att_lse.stride(0),
+            self.att_lse.stride(1),
+            self.att_lse.stride(2),
+            kv_group_num=self.group,
+            N_SPLITS=self.n_splits,
+            L_EXT=self.l_pad,
+            BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
+            BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
+            BLOCK_N=self.block_n,
+            MIN_BLOCK_KV=_MIN_BLOCK_KV,
+            num_warps=self.num_warps,
+            num_stages=1,
+            **_AMD_LAUNCH_KWARGS,
+        )
+
+    def _run_combine_kernel(
+        self, bs, q_extend, k_extend, v_extend, o_out, qo_indptr, sm_scale
+    ):
+        grid = (bs, self.h_q)
+        _verify_combine_stage2[grid](
+            self.att_out,
+            self.att_lse,
+            q_extend,
+            k_extend,
+            v_extend,
+            o_out,
+            sm_scale,
+            qo_indptr,
+            self.att_out.stride(0),
+            self.att_out.stride(1),
+            self.att_out.stride(2),
+            self.att_out.stride(3),
+            self.att_lse.stride(0),
+            self.att_lse.stride(1),
+            self.att_lse.stride(2),
+            q_extend.stride(0),
+            q_extend.stride(1),
+            k_extend.stride(0),
+            k_extend.stride(1),
+            v_extend.stride(0),
+            v_extend.stride(1),
+            o_out.stride(0),
+            o_out.stride(1),
+            kv_group_num=self.group,
+            N_SPLITS=self.n_splits,
+            L_EXT=self.l_pad,
+            BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
+            BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
+            num_warps=1,
+            num_stages=1,
+        )
+
+    def __call__(
+        self,
+        q_extend,
+        k_extend,
+        v_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        sm_scale,
+        o_out=None,
+        k_scale=1.0,
+        v_scale=1.0,
+    ):
+        if o_out is None:
+            o_out = torch.empty(
+                (q_extend.shape[0], self.h_q, self.v_head_dim),
+                dtype=q_extend.dtype,
+                device=q_extend.device,
+            )
+        # actual batch size for this call (<= max_bs); the grid uses it while the
+        # scratch buffers stay max_bs-sized (only the first bs slices are touched).
+        bs = qo_indptr.shape[0] - 1
+        # 1. prefix split-KV
+        self._run_prefix_kernel(
+            bs,
+            q_extend,
+            k_buffer,
+            v_buffer,
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            sm_scale,
+            k_scale,
+            v_scale,
+        )
+        # 2+3+4. fused combine + draft-draft + merge
+        self._run_combine_kernel(
+            bs,
+            q_extend,
+            k_extend,
+            v_extend,
+            o_out,
+            qo_indptr,
+            sm_scale,
+        )
+        return o_out
+
+
+# ---------------------------------------------------------------------------
+# Live-server dispatch entry.
+# ---------------------------------------------------------------------------
+# Cache one VerifySplitKV instance per (h_q, h_kv, head_dim, v_head_dim, l_ext,
+# device, n_splits) shape -- NOT keyed on the dynamic batch size. Buffers are
+# sized by the stable max_bs (grown only if a larger one is ever requested), so
+# a single instance serves every batch size: addresses stay fixed (graph-safe)
+# and GPU memory does not grow per batch size.
+_VK_CACHE = {}
+
+
+def _get_vk(
+    max_bs, h_q, h_kv, head_dim, v_head_dim, l_ext, device, n_splits=DEFAULT_N_SPLITS
+):
+    key = (h_q, h_kv, head_dim, v_head_dim, l_ext, str(device), n_splits)
+    vk = _VK_CACHE.get(key)
+    if vk is None:
+        block_n, num_warps = block_config(head_dim)
+        vk = VerifySplitKV(
+            max_bs,
+            h_q,
+            h_kv,
+            head_dim,
+            v_head_dim,
+            l_ext,
+            device=device,
+            n_splits=n_splits,
+            block_n=block_n,
+            num_warps=num_warps,
+        )
+        _VK_CACHE[key] = vk
+    else:
+        vk.grow_buffers(max_bs)
+    return vk
+
+
+def can_handle(
+    q_extend,
+    k_extend,
+    v_extend,
+    k_buffer,
+    v_buffer,
+    qo_indptr,
+    kv_indptr,
+    kv_indices,
+    custom_mask,
+    is_causal,
+    mask_indptr,
+    max_len_extend,
+    sliding_window_size=-1,
+    sinks=None,
+    logit_cap=0.0,
+    xai_temperature_len=-1,
+):
+    """Return True iff the split-KV verify path can serve this exact problem
+    with the same result as extend_attention_fwd. Conservative: anything not
+    explicitly handled -> False -> caller falls back to the baseline.
+
+    IMPORTANT: ``custom_mask`` is intentionally NOT inspected (its values can't
+    be read inside a captured HIP graph without a host sync). The kernel always
+    computes pure-causal attention, which equals the tree mask ONLY at
+    speculative topk == 1. The caller therefore MUST gate enablement on topk == 1
+    (TritonAttnBackend does: ``use_verify_splitkv = ... and self.topk == 1``).
+    At topk > 1 the tree is not causal and this path must stay disabled."""
+    # No exotic features.
+    if sinks is not None:
+        return False
+    if sliding_window_size is not None and sliding_window_size > 0:
+        return False
+    if logit_cap and logit_cap > 0:
+        return False
+    if xai_temperature_len is not None and xai_temperature_len > 0:
+        return False
+    if not is_causal:
+        return False
+    # q layout must be [tokens, H_Q, D]; head dims handled by power-of-2 pad.
+    if q_extend.dim() != 3 or k_extend.dim() != 3 or v_extend.dim() != 3:
+        return False
+    # GQA group must divide evenly.
+    h_q = q_extend.shape[1]
+    h_kv = k_extend.shape[1]
+    if h_kv == 0 or h_q % h_kv != 0:
+        return False
+    # head dims must match buffers.
+    if k_buffer.shape[1] != h_kv or v_buffer.shape[1] != h_kv:
+        return False
+    if q_extend.shape[2] != k_extend.shape[2]:
+        return False
+    if q_extend.shape[2] != k_buffer.shape[2]:
+        return False
+    if v_extend.shape[2] != v_buffer.shape[2]:
+        return False
+    # NOTE: must NOT read any tensor *values* here (no .item()/.cpu()): the
+    # target-verify step runs inside a captured CUDA/HIP graph, where a
+    # device->host sync raises hipErrorStreamCaptureUnsupported. We therefore
+    # gate purely on static shapes/dtypes/python scalars.
+    bs = qo_indptr.shape[0] - 1
+    if bs < 1:
+        return False
+    # max_len_extend must be a known positive python int (it is the static
+    # server_args.speculative_num_draft_tokens for the verify path). For
+    # topk=1 the per-seq extend len is constant == num_draft_tokens ==
+    # max_len_extend by construction of qo_indptr (arange with that step), so
+    # the L_EXT row-tile mask is exactly right and the tree custom_mask equals
+    # causal -- no value inspection required.
+    try:
+        mle = int(max_len_extend)
+    except (TypeError, ValueError):
+        return False
+    if mle < 1:
+        return False
+    # The packed extend tensor must hold exactly bs * max_len_extend rows
+    # (constant extend len). This is a pure shape check (no sync) and rejects
+    # any ragged/variable-extend batch -> falls back to the baseline.
+    if q_extend.shape[0] != bs * mle:
+        return False
+    return True
+
+
+def verify_splitkv_fwd(
+    q_extend,
+    k_extend,
+    v_extend,
+    o_extend,
+    k_buffer,
+    v_buffer,
+    qo_indptr,
+    kv_indptr,
+    kv_indices,
+    custom_mask,
+    is_causal,
+    mask_indptr,
+    max_len_extend,
+    k_scale,
+    v_scale,
+    sm_scale=None,
+    logit_cap=0.0,
+    skip_prefix_custom_mask=True,
+    sliding_window_size=-1,
+    sinks=None,
+    window_kv_offsets=None,
+    xai_temperature_len=-1,
+    max_bs=None,
+):
+    """Drop-in for extend_attention_fwd on the EAGLE target-verify (topk=1)
+    shape. Returns True if it ran (o_extend written), False if the case is
+    unsupported and the caller must fall back to extend_attention_fwd.
+
+    ``max_bs`` (optional) is the stable maximum batch size used to size the
+    cached scratch buffers; the backend passes its req_to_token_pool size. If
+    omitted it defaults to this call's bs.
+
+    Arg order mirrors extend_attention_fwd exactly so the call site is a
+    one-line swap.
+    """
+    if not can_handle(
+        q_extend,
+        k_extend,
+        v_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        custom_mask,
+        is_causal,
+        mask_indptr,
+        max_len_extend,
+        sliding_window_size=sliding_window_size,
+        sinks=sinks,
+        logit_cap=logit_cap,
+        xai_temperature_len=xai_temperature_len,
+    ):
+        return False
+
+    bs = qo_indptr.shape[0] - 1
+    h_q = q_extend.shape[1]
+    h_kv = k_extend.shape[1]
+    head_dim = q_extend.shape[2]
+    v_head_dim = v_extend.shape[2]
+    l_ext = int(max_len_extend)
+
+    if sm_scale is None:
+        sm_scale = 1.0 / (head_dim**0.5)
+    # k_scale/v_scale may be float or 0-d tensor; coerce to python float.
+    try:
+        k_scale = float(k_scale)
+    except (TypeError, ValueError):
+        k_scale = 1.0
+    try:
+        v_scale = float(v_scale)
+    except (TypeError, ValueError):
+        v_scale = 1.0
+
+    # Adaptive split count from the average prefix length. This is a
+    # pure-shape derivation (kv_indices.shape[0] / bs) -- no device->host sync,
+    # so it is safe inside a captured HIP graph. The whole batch shares one
+    # N_SPLITS (the grid dim must be a launch constexpr); the per-split kernel
+    # logic still clamps each split's [start,end) to that seq's real length, so
+    # mixed-length batches stay correct -- shorter seqs simply write fewer
+    # active splits (the rest emit the -inf lse sentinel, ignored in stage2).
+    avg_seqlen = kv_indices.shape[0] / max(1, bs)
+    n_splits = choose_n_splits(avg_seqlen)
+
+    # Size scratch by the stable max_bs (backend passes req_to_token_pool size);
+    # fall back to this call's bs if not provided / smaller.
+    if max_bs is None or max_bs < bs:
+        max_bs = bs
+    vk = _get_vk(
+        max_bs,
+        h_q,
+        h_kv,
+        head_dim,
+        v_head_dim,
+        l_ext,
+        q_extend.device,
+        n_splits=n_splits,
+    )
+    vk(
+        q_extend,
+        k_extend.contiguous(),
+        v_extend.contiguous(),
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        sm_scale,
+        o_out=o_extend,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+
+    return True

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__ = []

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",
+]

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

python/sglang/kernels/ops/gemm/chunked_embedding_lora_a.py

@@ -0,0 +1,142 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _chunked_embedding_lora_a_kernel(
+    # Pointers to tensors
+    input_ids,
+    weights,
+    output,
+    # Dimensions
+    vocab_size,
+    rank,
+    num_loras,
+    # Strides
+    w_stride_0,  # stride for lora index
+    w_stride_1,  # stride for rank
+    w_stride_2,  # stride for vocab
+    output_stride_0,
+    output_stride_1,
+    # Chunk info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    num_segments,
+    permutation,
+    # Meta-parameters
+    BLOCK_RANK: tl.constexpr,
+):
+    """
+    Embedding lookup for LoRA A weights without support for extra tokens.
+
+    Each program handles one chunk of tokens across rank dimension
+    """
+    chunk_idx = tl.program_id(axis=0)
+    # If chunk id is larger than actual number of chunks, skip
+    if chunk_idx >= num_segments:
+        return
+    chunk_start = tl.load(seg_indptr + chunk_idx)
+    chunk_end = tl.load(seg_indptr + chunk_idx + 1)
+    if chunk_start == chunk_end:
+        return
+    # Load LoRA adapter index for this segment, then look up the rank
+    lora_index = tl.load(weight_indices + chunk_idx)
+    rank_val = tl.load(lora_ranks + lora_index)
+    # If rank is 0, skip
+    if rank_val == 0:
+        return
+    # for each token in chunk, load embedding across rank dimension
+    for c in range(chunk_start, chunk_end):
+        s_index = tl.load(permutation + c)
+        # Load the token ID
+        token_id = tl.load(input_ids + s_index)
+        # Process in chunks of BLOCK_RANK dimensions
+        num_blocks = tl.cdiv(rank_val, BLOCK_RANK)
+
+        for block_id in range(num_blocks):
+            rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK
+            rank_mask = rank_offset < rank_val
+
+            # Use regular LoRA A weights
+            # weights shape: (num_loras, rank, vocab_size)
+            # We need to load weights[lora_index, rank_offset, token_id]
+            weight_ptr = (
+                weights
+                + lora_index * w_stride_0
+                + rank_offset * w_stride_1
+                + token_id * w_stride_2
+            )
+            emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0)
+
+            # Write to output
+            output_ptr = (
+                output + s_index * output_stride_0 + rank_offset * output_stride_1
+            )
+            tl.store(output_ptr, emb_values, mask=rank_mask)
+
+
+def chunked_embedding_lora_a_forward(
+    input_ids: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    vocab_size: int,
+) -> torch.Tensor:
+    """
+    Chunked Forward pass for LoRA A embedding lookup; each program handles one chunk of embedding lookup work
+    belonging to the same adapter
+
+    Args:
+        input_ids: (s,) token IDs
+        weights: (num_loras, rank, vocab_size) LoRA A embedding weights
+        batch_info: LoRABatchInfo containing batch information
+        vocab_size: base vocabulary size
+
+    Returns:
+        output: (s, rank) embedded features
+    """
+    assert input_ids.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(input_ids.shape) == 1
+    assert len(weights.shape) == 3
+
+    S = input_ids.shape[0]
+    num_loras = weights.shape[0]
+    rank = weights.shape[1]
+
+    # Block size for rank dimension
+    BLOCK_RANK = 128
+    num_segments = batch_info.num_segments
+    segment_grid = (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+    # 1D Grid: one program per chunk of embedding lookup work
+    grid = (segment_grid,)
+    output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype)
+
+    _chunked_embedding_lora_a_kernel[grid](
+        input_ids,
+        weights,
+        output,
+        vocab_size,
+        rank,
+        num_loras,
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        segment_grid,
+        batch_info.permutation,
+        BLOCK_RANK,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/chunked_sgmv_expand.py

@@ -0,0 +1,238 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+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
+
+
+@cached_triton_kernel(
+    lambda _, kwargs: (kwargs["NUM_SLICES"], kwargs["BLOCK_M"], kwargs["OUTPUT_DIM"])
+)
+@triton.jit(do_not_specialize=["num_segs", "output_stride_0", "output_stride_1"])
+def _chunked_lora_expand_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Output strides may differ from OUTPUT_DIM when compact LoRA output is
+    # accumulated into a wider base projection.
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    permutation,
+    num_segs,
+    # For fused output scaling
+    scalings,
+    # Offsets of q/k/v slice on output dimension
+    slice_offsets,
+    # Meta parameters
+    NUM_SLICES: tl.constexpr,
+    OUTPUT_DIM: tl.constexpr,
+    MAX_RANK: tl.constexpr,  # K = R
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """
+    Computes a chunked SGMV for LoRA expand operations.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (Tensor): The input tensor, which is the result of the LoRA A projection.
+            Shape: (s, num_slices * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, output_dim, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, output_dim) or a wider base output.
+    """
+    x_stride_0: tl.constexpr = NUM_SLICES * MAX_RANK
+    x_stride_1: tl.constexpr = 1
+
+    w_stride_0: tl.constexpr = OUTPUT_DIM * MAX_RANK
+    w_stride_1: tl.constexpr = MAX_RANK
+    w_stride_2: tl.constexpr = 1
+
+    pid_s = tl.program_id(axis=2)
+    if pid_s >= num_segs:
+        return
+
+    seg_start = tl.load(seg_indptr + pid_s)
+    seg_end = tl.load(seg_indptr + pid_s + 1)
+    if seg_start == seg_end:
+        return
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v)
+    w_index = tl.load(weight_indices + pid_s)
+    cur_rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if cur_rank == 0:
+        return
+
+    slice_id = tl.program_id(axis=1)
+    slice_start = tl.load(slice_offsets + slice_id)
+    slice_end = tl.load(slice_offsets + slice_id + 1)
+
+    scaling = tl.load(scalings + w_index)
+    # Adjust K (rank) according to the specific LoRA adapter
+    cur_rank = tl.minimum(MAX_RANK, cur_rank)
+
+    # Map logical sequence index to physical index
+    s_offset_logical = tl.arange(0, BLOCK_M) + seg_start
+    s_offset_physical = tl.load(
+        permutation + s_offset_logical, mask=s_offset_logical < seg_end, other=0
+    )
+
+    # Create pointers for the first block of x and weights[batch_id][n_start: n_end][:]
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    pid_n = tl.program_id(axis=0)
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + slice_start
+    k_offset = tl.arange(0, BLOCK_K)
+
+    x_ptrs = (
+        x
+        + slice_id * cur_rank * x_stride_1
+        + (s_offset_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    )
+    w_ptrs = (weights + w_index * w_stride_0) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(cur_rank, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset_logical[:, None] < seg_end)
+            & (k_offset[None, :] < cur_rank - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < cur_rank - k * BLOCK_K)
+            & (n_offset[None, :] < slice_end),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = output + (
+        s_offset_physical[:, None] * output_stride_0
+        + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset_logical[:, None] < seg_end) & (
+        n_offset[None, :] < slice_end
+    )
+    partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def chunked_sgmv_lora_expand_forward(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    slice_offsets: torch.Tensor,
+    max_slice_size: int,
+    base_output: Optional[torch.Tensor],
+) -> torch.Tensor:
+
+    # x: (s, slice_num * r)
+    # weights: (num_lora, output_dim, r)
+    # slice_offsets: boundaries for different slices in the output dimension
+    # output: (s, output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # For each slice i, accumulates:
+    # lora_output[:, slice_offsets[i]:slice_offsets[i+1]] += scaling * sgemm(x[:, i*cur_rank:(i+1)*cur_rank], weights[:, slice_offsets[i]:slice_offsets[i+1], :])
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    # Get dims
+    M = x.shape[0]
+    input_dim = x.shape[1]
+    OUTPUT_DIM = weights.shape[1]
+    MAX_RANK = weights.shape[2]
+    num_slices = len(slice_offsets) - 1
+    assert input_dim == num_slices * MAX_RANK
+
+    # Block shapes — use auto-tuned config if available, else defaults
+    BLOCK_M = batch_info.max_len
+    config = get_lora_expand_config(
+        K=OUTPUT_DIM, R=MAX_RANK, num_slices=num_slices, chunk_size=BLOCK_M
+    )
+    BLOCK_K = config["BLOCK_K"]
+    BLOCK_N = config["BLOCK_N"]
+
+    num_segments = batch_info.num_segments
+    segment_grid = (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+
+    grid = (
+        triton.cdiv(max_slice_size, BLOCK_N),
+        num_slices,  # number of slices in the input/output
+        segment_grid,
+    )
+
+    if base_output is None:
+        output = torch.zeros((M, OUTPUT_DIM), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    # Optional launch params from tuned config
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+    if "maxnreg" in config:
+        extra_kwargs["maxnreg"] = config["maxnreg"]
+
+    _chunked_lora_expand_kernel[grid](
+        x=x,
+        weights=weights,
+        output=output,
+        output_stride_0=output.stride(0),
+        output_stride_1=output.stride(1),
+        seg_indptr=batch_info.seg_indptr,
+        weight_indices=batch_info.weight_indices,
+        lora_ranks=batch_info.lora_ranks,
+        permutation=batch_info.permutation,
+        num_segs=segment_grid,
+        scalings=batch_info.scalings,
+        slice_offsets=slice_offsets,
+        # constants
+        NUM_SLICES=num_slices,
+        OUTPUT_DIM=OUTPUT_DIM,
+        MAX_RANK=MAX_RANK,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        **extra_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/chunked_sgmv_shrink.py

@@ -0,0 +1,196 @@
+import torch
+import triton
+import triton.language as tl
+
+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
+
+
+@cached_triton_kernel(
+    lambda _, kwargs: (kwargs["K"], kwargs["NUM_SLICES"], kwargs["BLOCK_M"])
+)
+@triton.jit(do_not_specialize=["num_segs"])
+def _chunked_lora_shrink_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Information on sequence lengths,ranks and weight id
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    permutation,
+    num_segs,
+    # Meta parameters
+    N: tl.constexpr,  # num_slices * r
+    K: tl.constexpr,  # input_dim
+    NUM_SLICES: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """
+    Computes a chunked SGMV for LoRA shrink operations.
+
+    The kernel ensures that output[seg_start:seg_start + seg_len, :rank * num_slices]
+    stores the product of the input `x` and the LoRA weights for the corresponding
+    sequence. This implies that when rank is 0, the kernel is essentially a no-op,
+    as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty).
+
+    Args:
+        x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s`
+            is the sum of all sequence lengths in the batch.
+        weights (torch.Tensor): The LoRA A weights for all available adapters,
+            with shape `(num_lora, N, K)` where N = num_slices * r.
+        output (torch.Tensor): The output tensor of shape `(s, N)`.
+    """
+    x_stride_1: tl.constexpr = 1
+    x_stride_0: tl.constexpr = K
+
+    w_stride_0: tl.constexpr = N * K
+    w_stride_1: tl.constexpr = K
+    w_stride_2: tl.constexpr = 1
+
+    output_stride_0: tl.constexpr = N
+    output_stride_1: tl.constexpr = 1
+
+    pid_s = tl.program_id(1)
+    if pid_s >= num_segs:
+        return
+
+    pid_n = tl.program_id(0)
+
+    seg_start = tl.load(seg_indptr + pid_s)
+    seg_end = tl.load(seg_indptr + pid_s + 1)
+    if seg_start == seg_end:
+        return
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    w_index = tl.load(weight_indices + pid_s)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel becomes a no-op as the output is always trivially correct.
+    if rank == 0:
+        return
+
+    # Adjust N dim according to the specific LoRA adapter
+    cur_n = tl.minimum(N, rank * NUM_SLICES)
+
+    # Map logical sequence index to physical index
+    s_offset_logical = tl.arange(0, BLOCK_M) + seg_start
+    s_offset_physical = tl.load(
+        permutation + s_offset_logical, mask=s_offset_logical < seg_end, other=0
+    )
+
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+    x_ptrs = x + (
+        s_offset_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1
+    )
+    w_ptrs = (weights + w_index * w_stride_0) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset_logical[:, None] < seg_end)
+            & (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < cur_n),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = output + (
+        s_offset_physical[:, None] * output_stride_0
+        + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset_logical[:, None] < seg_end) & (n_offset[None, :] < cur_n)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def chunked_sgmv_lora_shrink_forward(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    num_slices: int,
+) -> torch.Tensor:
+    # x: (s, input_dim)
+    # weights: (num_lora, num_slices * r, input_dim)
+    # output: (s, num_slices * r)
+    # num_slices: qkv=3, gate_up=2, others=1
+    # when called with multiple slices, the weights.shape[-2] will be num_slices * r
+    # input_dim is much larger than r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    # Block shapes — use auto-tuned config if available, else defaults
+    BLOCK_M = batch_info.max_len
+    # weights shape is (num_lora, num_slices * rank, input_dim)
+    MAX_RANK = weights.shape[1] // num_slices
+    config = get_lora_shrink_config(
+        K=weights.shape[2], R=MAX_RANK, num_slices=num_slices, chunk_size=BLOCK_M
+    )
+    BLOCK_N = config["BLOCK_N"]
+    BLOCK_K = config["BLOCK_K"]
+
+    S = x.shape[0]
+    N = weights.shape[1]
+    K = weights.shape[2]
+    assert x.shape[-1] == K
+
+    num_segments = batch_info.num_segments
+    segment_grid = (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+    grid = (
+        triton.cdiv(N, BLOCK_N),
+        segment_grid,
+    )
+
+    # Optional launch params from tuned config
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+
+    output = torch.empty((S, N), device=x.device, dtype=x.dtype)
+    _chunked_lora_shrink_kernel[grid](
+        x=x,
+        weights=weights,
+        output=output,
+        seg_indptr=batch_info.seg_indptr,
+        weight_indices=batch_info.weight_indices,
+        lora_ranks=batch_info.lora_ranks,
+        permutation=batch_info.permutation,
+        num_segs=segment_grid,
+        # constants
+        N=N,
+        K=K,
+        NUM_SLICES=num_slices,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        **extra_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 1,
+        "maxnreg": 128
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 2,
+        "maxnreg": 128
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3,
+        "maxnreg": 160
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 1,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 2,
+        "maxnreg": 128
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "maxnreg": 112
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 3,
+        "maxnreg": 128
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 128,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 4
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 128,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    }
+}

python/sglang/kernels/ops/gemm/embedding_lora_a.py

@@ -0,0 +1,186 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _embedding_lora_a_kernel(
+    # Pointers to tensors
+    input_ids,
+    weights,
+    output,
+    extra_embeddings,
+    # Dimensions
+    vocab_size,
+    rank,
+    num_loras,
+    # Strides
+    w_stride_0,  # stride for lora index
+    w_stride_1,  # stride for rank
+    w_stride_2,  # stride for vocab
+    output_stride_0,
+    output_stride_1,
+    extra_emb_stride_0,  # stride for lora index
+    extra_emb_stride_1,  # stride for token
+    extra_emb_stride_2,  # stride for hidden dim (= rank for extra embeddings)
+    # Batch info
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    # Meta-parameters
+    BLOCK_RANK: tl.constexpr,
+    HAS_EXTRA_EMBEDDINGS: tl.constexpr,
+):
+    """
+    Embedding lookup for LoRA A weights with support for extra tokens.
+
+    Each program handles one token across a block of rank dimensions.
+    Grid: (cdiv(max_len, 1), bs) - one program per token in each batch
+    """
+    batch_id = tl.program_id(axis=1)
+    token_idx = tl.program_id(axis=0)
+
+    w_index = tl.load(weight_indices + batch_id)
+    rank_val = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, skip
+    if rank_val == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_len = tl.load(seg_lens + batch_id)
+
+    # Check if this token is within the segment
+    if token_idx >= seg_len:
+        return
+
+    # Load the token ID
+    token_id = tl.load(input_ids + seg_start + token_idx)
+
+    # Process in chunks of BLOCK_RANK dimensions
+    num_blocks = tl.cdiv(rank_val, BLOCK_RANK)
+
+    for block_id in range(num_blocks):
+        rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK
+        rank_mask = rank_offset < rank_val
+
+        # Check if this is an extra token
+        is_extra_token = token_id >= vocab_size
+
+        if HAS_EXTRA_EMBEDDINGS and is_extra_token:
+            # Use extra embeddings
+            extra_token_id = token_id - vocab_size
+            extra_emb_ptr = (
+                extra_embeddings
+                + w_index * extra_emb_stride_0
+                + extra_token_id * extra_emb_stride_1
+                + rank_offset * extra_emb_stride_2
+            )
+            emb_values = tl.load(extra_emb_ptr, mask=rank_mask, other=0.0)
+        else:
+            # Use regular LoRA A weights
+            # weights shape: (num_loras, rank, vocab_size)
+            # We need to load weights[w_index, rank_offset, token_id]
+            token_id_clamped = tl.minimum(token_id, vocab_size - 1)
+            weight_ptr = (
+                weights
+                + w_index * w_stride_0
+                + rank_offset * w_stride_1
+                + token_id_clamped * w_stride_2
+            )
+            emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0)
+
+        # Write to output
+        output_ptr = (
+            output
+            + (seg_start + token_idx) * output_stride_0
+            + rank_offset * output_stride_1
+        )
+        tl.store(output_ptr, emb_values, mask=rank_mask)
+
+
+def embedding_lora_a_fwd(
+    input_ids: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    vocab_size: int,
+    extra_embeddings: torch.Tensor = None,
+) -> torch.Tensor:
+    """
+    Forward pass for LoRA A embedding lookup.
+
+    Args:
+        input_ids: (s,) token IDs
+        weights: (num_loras, rank, vocab_size) LoRA A embedding weights
+        batch_info: LoRABatchInfo containing batch information
+        vocab_size: base vocabulary size
+        extra_embeddings: (num_loras, num_extra_tokens, rank) extra token embeddings
+
+    Returns:
+        output: (s, rank) embedded features
+    """
+    assert input_ids.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(input_ids.shape) == 1
+    assert len(weights.shape) == 3
+
+    S = input_ids.shape[0]
+    num_loras = weights.shape[0]
+    rank = weights.shape[1]
+    vocab_size_weights = weights.shape[2]
+
+    # Block size for rank dimension
+    BLOCK_RANK = 128
+
+    has_extra_embeddings = extra_embeddings is not None
+
+    if has_extra_embeddings:
+        assert extra_embeddings.is_contiguous()
+        extra_emb_stride = (
+            extra_embeddings.stride(0),
+            extra_embeddings.stride(1),
+            extra_embeddings.stride(2),
+        )
+    else:
+        # Create dummy tensor to satisfy Triton
+        extra_embeddings = torch.empty(
+            (1, 1, 1), device=input_ids.device, dtype=weights.dtype
+        )
+        extra_emb_stride = (1, 1, 1)
+
+    # Grid: one program per token in each batch segment
+    grid = (
+        batch_info.max_len,
+        batch_info.bs,
+    )
+
+    output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype)
+
+    _embedding_lora_a_kernel[grid](
+        input_ids,
+        weights,
+        output,
+        extra_embeddings,
+        vocab_size,
+        rank,
+        num_loras,
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        extra_emb_stride[0],
+        extra_emb_stride[1],
+        extra_emb_stride[2],
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        BLOCK_RANK,
+        has_extra_embeddings,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/gate_up_lora_b.py

@@ -0,0 +1,204 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _gate_up_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    output_dim,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling
+    scalings,
+):
+    """
+    This kernel packs 2 sgemms (gate/up) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (Tensor): The input tensor, which is the result of the LoRA A projection.
+            Shape: (s, 2 * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, 2 * output_dim, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, 2 * output_dim).
+    """
+    # output_dim >> K
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # gate_up_id decides which of gate or up (0: gate, 1: up)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    gate_up_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+    seg_len = tl.load(seg_lens + batch_id)
+    if seg_len == 0:
+        return
+    seg_start = tl.load(seg_indptr + batch_id)
+    n_start = gate_up_id * output_dim  # offset on output dim
+    scaling = tl.load(scalings + w_index)
+
+    # Adjust K (rank) according to the specific LoRA adapter
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(output_dim, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    # Create pointers for the first block of x and weights
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = (
+        x
+        + (gate_up_id * K) * x_stride_1
+        + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    )
+    w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K)
+            & (n_offset[None, :] < output_dim),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = (
+        output
+        + n_start * output_stride_1
+        + (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1)
+    )
+    output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < output_dim)
+    partial_sum += tl.load(output_ptr, mask=output_mask)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def gate_up_lora_b_fwd(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+
+    # x: (s, 2 * r)
+    # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+    # output: (s, 2 * output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :])
+    # lora_output[:, output_dim:]
+    #      = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = gate_up_lora_b.shape[-1]
+    assert input_dim == 2 * r
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT),
+        2,  # this dimension decides current block computes on gate or up proj
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, 2 * output_dim), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    _gate_up_lora_b_kernel[grid_b](
+        x,
+        gate_up_lora_b,
+        output,
+        r,
+        output_dim,
+        x.stride(0),
+        x.stride(1),
+        gate_up_lora_b.stride(0),
+        gate_up_lora_b.stride(1),
+        gate_up_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/kernel_utils.py

@@ -0,0 +1,19 @@
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _resolve_token_positions(
+    sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER: tl.constexpr
+):
+    """Map logical segment offsets to physical token positions.
+
+    When SORTED_BY_ADAPTER is True, segments are grouped by adapter and
+    sorted_token_ids provides the indirection to the original token rows.
+    When False, tokens are already contiguous starting at seg_start.
+    """
+    if SORTED_BY_ADAPTER:
+        return tl.load(
+            sorted_token_ids + seg_start + s_offset, mask=s_offset < seg_len
+        ).to(tl.int64)
+    return (seg_start + s_offset).to(tl.int64)

python/sglang/kernels/ops/gemm/kv_b_lora_absorbed.py

@@ -0,0 +1,853 @@
+"""Triton kernels for absorbed-MLA ``kv_b_proj`` LoRA correction.
+
+The absorbed-MLA path bypasses ``kv_b_proj.forward()`` and folds the K/V
+sides as plain BMMs ``q_nope @ w_kc`` and ``attn_output @ w_vc``.  When a
+LoRA adapter is active on ``kv_b_proj`` we add the LoRA delta to
+``q_nope_out`` / ``attn_bmm_output`` manually.
+
+Using the standard LoRA factored math we *never* materialize ``B @ A``:
+
+    q_correction = q_nope     @ B_kc @ A   * scaling          # K-side
+    v_correction = attn_output @ A.T @ B_vc.T * scaling       # V-side
+
+where ``A: (slot, rank, kv_lora_rank)`` is the LoRA-A of ``kv_b_proj``
+(shared across heads) and ``B: (slot, num_heads*(qk_nope+v_head_dim), rank)``
+is the LoRA-B; ``B_kc`` / ``B_vc`` are its K-half / V-half slices.
+
+Four kernels split the math along the factorization boundary, all using
+the SGMM idiom from ``sgemm_lora_a`` / ``qkv_lora_b`` and the segment-indptr
+routing used by ``chunked_sgmv_*``:
+
+  * ``step_a_q_fwd``: per-head per-slot SGMM, ``(S,H,qk_nope) -> (S,H,rank)``
+  * ``step_b_q_fwd``: shared-A per-slot SGMM, scaled+accumulated,
+    ``(S,H,rank) -> (S,H,kv_lora_rank)``
+  * ``step_a_v_fwd``: shared-A.T per-slot SGMM, ``(S,H,kv_lora_rank) -> (S,H,rank)``
+  * ``step_b_v_fwd``: per-head per-slot SGMM with V-half of B, transposed,
+    scaled+accumulated, ``(S,H,rank) -> (S,H,v_head_dim)``
+
+Grid axes for each kernel:
+  axis 0 : output tile in (S, N)         -- tile_id = pid_s * num_pid_n + pid_n
+  axis 1 : head_id                       -- per-head weight slice
+  axis 2 : batch_id (segment / request)  -- per-slot weight routing via weight_indices
+
+Per-segment routing: each program derives its segment length from
+``seg_indptr[segment_id + 1] - seg_indptr[segment_id]``, loads
+``weight_indices[segment_id]`` once, and uses that slot's slice of the LoRA
+weight stack.  When ``permutation`` is present, rows are routed through it,
+matching the csgmv backend's adapter-grouped chunks.  No Python loops over slots
+or heads.
+
+The math also stays in the input dtype (no fp32 round-trip) -- the
+contraction dim ``rank`` is small (typically 16-64), so bf16 accumulation
+over it is acceptable.  ``tl.dot`` itself uses fp32 accumulation internally.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+# ---------------------------------------------------------------------------
+# Block sizes -- chosen per-kernel from the natural shape of each step.
+#
+# The factored math gives the four kernels these contraction (K) and output
+# (N) ranges (for Kimi-K2.5: rank=16-32, qk_nope=v_head_dim=128, kv_lora_rank=512):
+#
+#                              K (contraction)      N (output)
+#   step_a_q                   qk_nope (~128)       rank (~16-32)
+#   step_b_q                   rank (~16-32)        kv_lora_rank (~512)
+#   step_a_v                   kv_lora_rank (~512)  rank (~16-32)
+#   step_b_v                   rank (~16-32)        v_head_dim (~128)
+#
+# So the "step_a_*" kernels want a large BLOCK_K (to keep loop iters small)
+# and a small BLOCK_N (matched to rank to avoid wasted tile lanes), while
+# the "step_b_*" kernels are the inverse.  Kernels aren't autotuned -- the
+# decode-shape workload is too small to benefit and the sweep surface is
+# wide.
+# ---------------------------------------------------------------------------
+
+_BLOCK_S = 16
+_STEP_A_BLOCK_K = 64  # contraction over qk_nope (~128) or kv_lora_rank (~512)
+_STEP_A_BLOCK_N = 16  # output is rank
+_STEP_B_BLOCK_K = 16  # contraction is rank
+_STEP_B_BLOCK_N = 64  # output is kv_lora_rank (~512) or v_head_dim (~128)
+
+
+def _num_segments(batch_info: LoRABatchInfo) -> int:
+    return batch_info.num_segments or batch_info.bs
+
+
+def _max_segment_len(batch_info: LoRABatchInfo) -> int:
+    if batch_info.max_len is not None:
+        return batch_info.max_len
+    if batch_info.seg_lens is not None:
+        return int(batch_info.seg_lens.max().item())
+    raise ValueError("LoRA batch_info must provide max_len or seg_lens.")
+
+
+def _segment_grid_size(batch_info: LoRABatchInfo, num_segments: int) -> int:
+    return (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+
+
+# ---------------------------------------------------------------------------
+# Kernel 1 -- Step A_q: per-head per-slot SGMM, reads K-half of B
+#
+#     q_lora_a[t, h, r] = sum_{i<qk_nope} q_nope[t, h, i] * B[slot, h*FULL_K + i, r]
+#
+# x      : (S, H, qk_nope)
+# w (B)  : (num_lora, H*FULL_K, rank)   -- FULL_K = qk_nope + v_head_dim
+# out    : (S, H, rank)                 -- fresh allocation, no accumulate
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_q_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    H_FULL_K,  # H * (qk_nope + v_head_dim), the row-stride landmark
+    K,  # qk_nope (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,
+    w_stride_k,
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # per-head row stride in B (qk_nope + v_head_dim)
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank (allows mixed-rank batches).
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    # Clamp masked-lane indices into the valid range so pointer arithmetic
+    # stays in-bounds even before the load mask drops the values.
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    head_row_base = (
+        head_id * FULL_K
+    )  # row offset for this head's K-half (i in [0, qk_nope))
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + i, r]: row dim of B carries i (= GEMM K),
+        # column dim carries r (= GEMM N).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + (head_row_base + safe_k[:, None]) * w_stride_n
+            + safe_n[None, :] * w_stride_k,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_q_fwd(
+    q_nope: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    full_K_per_head: int,
+) -> torch.Tensor:
+    """Step A of the q-side correction.
+
+    Args:
+        q_nope: ``(S, H, qk_nope)``, the absorbed-MLA q intermediate.
+        B_buf: ``(num_lora, H*full_K_per_head, rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        full_K_per_head: ``qk_nope + v_head_dim``, the row stride per head in B.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate, ready for step B_q.
+    """
+    S, H, qk_nope_dim = q_nope.shape
+    rank = B_buf.shape[-1]
+    out = torch.empty((S, H, rank), device=q_nope.device, dtype=q_nope.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, _STEP_A_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_a_q_kernel[grid](
+        q_nope,
+        B_buf,
+        out,
+        S,
+        H * full_K_per_head,
+        qk_nope_dim,
+        rank,
+        q_nope.stride(0),
+        q_nope.stride(1),
+        q_nope.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_A_BLOCK_N,
+        BLOCK_K=_STEP_A_BLOCK_K,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 2 -- Step B_q: shared-A per-slot SGMM, scaled + accumulated
+#
+#     base[t, h, k] += sum_r x[t, h, r] * A[slot, r, k] * scaling
+#
+# x      : (S, H, rank)
+# w (A)  : (num_lora, rank, kv_lora_rank)
+# base   : (S, H, kv_lora_rank), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_q_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # kv_lora_rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_k,
+    w_stride_n,
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    # Truncate contraction K to this slot's rank.
+    K_eff = tl.minimum(K, cur_rank)
+
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = tl.minimum(n_offset, N - 1)
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]  (k iterates over rank)
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, k, n]: read k along contraction, n along output.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    # Accumulate into base[s, h, n].
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    partial_sum += tl.load(base + base_offs, mask=out_mask, other=0.0)
+    tl.store(base + base_offs, partial_sum, mask=out_mask)
+
+
+def step_b_q_fwd(
+    q_lora_a: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Step B of the q-side correction, accumulating into ``base_output``.
+
+    Args:
+        q_lora_a: ``(S, H, rank)`` from step A_q.
+        A_buf: ``(num_lora, rank, kv_lora_rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, kv_lora_rank)``, modified in-place
+            (the absorbed ``q_nope @ w_kc`` result).
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = q_lora_a.shape
+    kv_lora_rank = A_buf.shape[-1]
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(kv_lora_rank, _STEP_B_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_b_q_kernel[grid](
+        q_lora_a,
+        A_buf,
+        base_output,
+        S,
+        rank,
+        kv_lora_rank,
+        q_lora_a.stride(0),
+        q_lora_a.stride(1),
+        q_lora_a.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+    )
+    return base_output
+
+
+# ---------------------------------------------------------------------------
+# Kernel 3 -- Step A_v: shared-A.T per-slot SGMM (no scaling, fresh output)
+#
+#     attn_lora_a[t, h, r] = sum_k attn_output[t, h, k] * A[slot, r, k]
+#
+# x      : (S, H, kv_lora_rank)
+# w (A)  : (num_lora, rank, kv_lora_rank) -- accessed transposed vs step B_q
+# out    : (S, H, rank), fresh allocation
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_v_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    K,  # kv_lora_rank (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # A's "rank" axis (= GEMM N)
+    w_stride_k,  # A's "kv_lora_rank" axis (= GEMM K)
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank.
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, r, k] -- here we want each (k, r) so we read along k
+        # (inner / contraction) and produce r as output.  Stride access:
+        # the row dim is r (= GEMM N), column dim is k (= GEMM K).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_v_fwd(
+    attn_output: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+) -> torch.Tensor:
+    """Step A of the v-side correction.
+
+    Args:
+        attn_output: ``(S, H, kv_lora_rank)``, the post-attention intermediate.
+        A_buf: ``(num_lora, rank, kv_lora_rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate for step B_v.
+    """
+    S, H, kv_lora_rank = attn_output.shape
+    rank = A_buf.shape[1]
+    out = torch.empty((S, H, rank), device=attn_output.device, dtype=attn_output.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, _STEP_A_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_a_v_kernel[grid](
+        attn_output,
+        A_buf,
+        out,
+        S,
+        kv_lora_rank,
+        rank,
+        attn_output.stride(0),
+        attn_output.stride(1),
+        attn_output.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_A_BLOCK_N,
+        BLOCK_K=_STEP_A_BLOCK_K,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 4 -- Step B_v: per-head per-slot SGMM with V-half of B (transposed),
+# scaled + accumulated
+#
+#     base[t, h, j] += sum_r x[t, h, r] * B[slot, h*FULL_K + qk_nope + j, r] * scaling
+#
+# x      : (S, H, rank)
+# w (B)  : (num_lora, H*FULL_K, rank), V-half slice via offset
+# base   : (S, H, v_head_dim), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_v_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # v_head_dim (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # B's row dim (h*FULL_K + j) -- this is GEMM N
+    w_stride_k,  # B's rank dim -- this is GEMM K
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # qk_nope + v_head_dim
+    QK_NOPE_OFFSET: tl.constexpr,  # offset of V-half within each head's row block
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    K_eff = tl.minimum(K, cur_rank)
+
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = tl.minimum(n_offset, N - 1)
+
+    # V-half row base for this head: h*FULL_K + qk_nope
+    head_row_base = head_id * FULL_K + QK_NOPE_OFFSET
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + qk_nope + j, r] -- row dim is j (= GEMM N),
+        # column dim is r (= GEMM K).  Transposed access vs step A_q.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + (head_row_base + safe_n[None, :]) * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    partial_sum += tl.load(base + base_offs, mask=out_mask, other=0.0)
+    tl.store(base + base_offs, partial_sum, mask=out_mask)
+
+
+def step_b_v_fwd(
+    attn_lora_a: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+    qk_nope_head_dim: int,
+    v_head_dim: int,
+) -> torch.Tensor:
+    """Step B of the v-side correction, accumulating into ``base_output``.
+
+    Args:
+        attn_lora_a: ``(S, H, rank)`` from step A_v.
+        B_buf: ``(num_lora, H*(qk_nope+v_head_dim), rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, v_head_dim)``, modified in-place
+            (the absorbed ``attn_output @ w_vc`` result).
+        qk_nope_head_dim: offset of V-half within each head's row block of B.
+        v_head_dim: output feature dim per head.
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = attn_lora_a.shape
+    full_K_per_head = qk_nope_head_dim + v_head_dim
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(v_head_dim, _STEP_B_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_b_v_kernel[grid](
+        attn_lora_a,
+        B_buf,
+        base_output,
+        S,
+        rank,
+        v_head_dim,
+        attn_lora_a.stride(0),
+        attn_lora_a.stride(1),
+        attn_lora_a.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        QK_NOPE_OFFSET=qk_nope_head_dim,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+    )
+    return base_output

python/sglang/kernels/ops/gemm/lora_tuning_config.py

@@ -0,0 +1,201 @@
+"""
+Configuration loader for auto-tuned LoRA CSGMV kernel block sizes.
+
+Follows the same pattern as fused_moe_triton_config.py:
+- Offline tuning script writes JSON files keyed by chunk_size (BLOCK_M)
+- At server startup, the config loader reads the best block sizes for each kernel
+- Kernels use these instead of hardcoded defaults
+
+Config file naming: lora_{kernel},K={K},R={R},S={S},device={device}.json
+Where kernel is "shrink" or "expand", K is input_dim, R is max_rank, S is num_slices.
+
+Config file format (keyed by chunk_size):
+{
+    "16": {"BLOCK_N": 16, "BLOCK_K": 256, "num_warps": 4, "num_stages": 3},
+    "32": {"BLOCK_N": 32, "BLOCK_K": 128, "num_warps": 4, "num_stages": 4},
+    "128": {"BLOCK_N": 64, "BLOCK_K": 256, "num_warps": 8, "num_stages": 3}
+}
+
+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/kernels/ops/gemm/configs/
+
+    # Server automatically picks them up:
+    python3 -m sglang.launch_server --model ... --enable-lora --lora-backend csgmv
+"""
+
+from __future__ import annotations
+
+import functools
+import json
+import logging
+import os
+from typing import Any, Dict, Optional
+
+import triton
+
+from sglang.srt.utils import get_device_name
+
+logger = logging.getLogger(__name__)
+
+
+def get_lora_config_file_name(
+    kernel: str,
+    K: int,
+    R: int,
+    S: int,
+) -> str:
+    """Generate config filename for a LoRA kernel configuration.
+
+    Args:
+        kernel: "shrink" or "expand"
+        K: The large dimension (input_dim for shrink, output_dim for expand)
+        R: The max LoRA rank
+        S: num_slices (qkv=3, gate_up=2, others=1)
+    """
+    device_name = get_device_name().replace(" ", "_")
+    return f"lora_{kernel},K={K},R={R},S={S},device={device_name}.json"
+
+
+@functools.lru_cache
+def get_lora_configs(
+    kernel: str,
+    K: int,
+    R: int,
+    S: int,
+) -> Optional[Dict[int, Dict[str, Any]]]:
+    """Load pre-tuned LoRA kernel configs from JSON files.
+
+    Returns a dict mapping chunk_size (BLOCK_M) to block size configs,
+    or None if no config file is found.
+    """
+    json_file_name = get_lora_config_file_name(kernel, K, R, S)
+
+    config_dir = os.environ.get(
+        "SGLANG_LORA_CONFIG_DIR", os.path.dirname(os.path.realpath(__file__))
+    )
+    configs_root = os.path.join(config_dir, "csgmv_configs")
+
+    triton_version = triton.__version__
+    version_dir = f"triton_{triton_version.replace('.', '_')}"
+
+    # Try exact triton version first
+    config_file_path = os.path.join(configs_root, version_dir, json_file_name)
+    if os.path.exists(config_file_path):
+        with open(config_file_path) as f:
+            logger.info(f"Using LoRA {kernel} config from {config_file_path}.")
+            return {int(key): val for key, val in json.load(f).items()}
+
+    # Scan existing version directories as fallback (newest first)
+    if os.path.isdir(configs_root):
+        version_dirs = sorted(
+            (d for d in os.listdir(configs_root) if d.startswith("triton_")),
+            reverse=True,
+        )
+        for vdir in version_dirs:
+            if vdir == version_dir:
+                continue
+            try_path = os.path.join(configs_root, vdir, json_file_name)
+            if os.path.exists(try_path):
+                with open(try_path) as f:
+                    logger.warning(
+                        f"LoRA {kernel} config not found for Triton {triton_version}. "
+                        f"Falling back to {try_path}."
+                    )
+                    return {int(key): val for key, val in json.load(f).items()}
+
+    return None
+
+
+# Default block sizes (current hardcoded values)
+DEFAULT_SHRINK_CONFIG = {"BLOCK_N": 16, "BLOCK_K": 256}
+DEFAULT_EXPAND_CONFIG = {"BLOCK_N": 64, "BLOCK_K": 16}
+
+# Track which configs have been logged to avoid spamming on every forward pass
+_logged_configs: set = set()
+
+
+def get_lora_shrink_config(
+    K: int,
+    R: int,
+    num_slices: int,
+    chunk_size: int,
+) -> Dict[str, int]:
+    """Get block sizes for the CSGMV shrink (lora_a) kernel.
+
+    Args:
+        K: input_dim
+        R: max_rank
+        num_slices: number of slices (qkv=3, gate_up=2, others=1)
+        chunk_size: BLOCK_M value (= batch_info.max_len)
+    """
+    log_key = ("shrink", K, R, num_slices, chunk_size)
+    configs = get_lora_configs("shrink", K, R, num_slices)
+    if configs is not None:
+        config = configs.get(chunk_size)
+        if config is None:
+            closest = min(configs.keys(), key=lambda x: abs(x - chunk_size))
+            config = configs[closest]
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA shrink (K={K}, R={R}): no config for chunk_size={chunk_size}, "
+                    f"using closest={closest}: {config}"
+                )
+        else:
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA shrink (K={K}, R={R}, chunk_size={chunk_size}): tuned config {config}"
+                )
+        return config
+    if log_key not in _logged_configs:
+        _logged_configs.add(log_key)
+        logger.info(
+            f"LoRA shrink (K={K}, R={R}): no tuned config, using defaults {DEFAULT_SHRINK_CONFIG}"
+        )
+    return dict(DEFAULT_SHRINK_CONFIG)
+
+
+def get_lora_expand_config(
+    K: int,
+    R: int,
+    num_slices: int,
+    chunk_size: int,
+) -> Dict[str, int]:
+    """Get block sizes for the CSGMV expand (lora_b) kernel.
+
+    Args:
+        K: output_dim
+        R: max_rank
+        num_slices: number of slices (qkv=3, gate_up=2, others=1)
+        chunk_size: BLOCK_M value (= batch_info.max_len)
+    """
+    log_key = ("expand", K, R, num_slices, chunk_size)
+    configs = get_lora_configs("expand", K, R, num_slices)
+    if configs is not None:
+        config = configs.get(chunk_size)
+        if config is None:
+            closest = min(configs.keys(), key=lambda x: abs(x - chunk_size))
+            config = configs[closest]
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA expand (K={K}, R={R}): no config for chunk_size={chunk_size}, "
+                    f"using closest={closest}: {config}"
+                )
+        else:
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA expand (K={K}, R={R}, chunk_size={chunk_size}): tuned config {config}"
+                )
+        return config
+    if log_key not in _logged_configs:
+        _logged_configs.add(log_key)
+        logger.info(
+            f"LoRA expand (K={K}, R={R}): no tuned config, using defaults {DEFAULT_EXPAND_CONFIG}"
+        )
+    return dict(DEFAULT_EXPAND_CONFIG)

python/sglang/kernels/ops/gemm/qkv_lora_b.py

@@ -0,0 +1,216 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _qkv_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    max_qkv_out_dim,  # max(output_q_dim, output_kv_dim)
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    # Offsets of q/k/v slice on output dimension
+    n_offs,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling
+    scalings,
+):
+    """
+    This kernel packs 3 sgemms (q/k/v) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (Tensor): The input tensor, which is the result of the LoRA A projection.
+            Shape: (s, 3 * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank. The second dimension is partitioned
+            for Q, K, and V.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, N_Q + 2 * N_KV, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, N_Q + 2 * N_KV).
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    qkv_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+    seg_len = tl.load(seg_lens + batch_id)
+    if seg_len == 0:
+        return
+    seg_start = tl.load(seg_indptr + batch_id)
+    n_start = tl.load(n_offs + qkv_id)
+    n_size = tl.load(n_offs + qkv_id + 1) - n_start
+    scaling = tl.load(scalings + w_index)
+    # Adjust K (rank) according to the specific LoRA adapter
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(max_qkv_out_dim, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    # Create pointers for the first block of x and weights[batch_id][n_start: n_end][:]
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = (
+        x
+        + (qkv_id * K) * x_stride_1
+        + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    )
+    w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < n_size),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = (
+        output
+        + n_start * output_stride_1
+        + (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1)
+    )
+    output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < n_size)
+    partial_sum += tl.load(output_ptr, mask=output_mask)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def qkv_lora_b_fwd(
+    x: torch.Tensor,
+    qkv_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_offset: torch.Tensor,
+    max_qkv_out_dim: int,
+    base_output: torch.Tensor = None,
+    n_slices: int = 3,
+) -> torch.Tensor:
+
+    # x: (s, n_slices * r)
+    # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
+    # output_offset = [0, output_dim_q, output_dim_q + output_dim_kv,
+    #                     output_dim_q + 2 * output_dim_kv]  (length n_slices + 1)
+    # max_qkv_out_dim = max(output_dim_q, output_dim_kv)
+    # output: (s, output_dim_q + 2 * output_dim_kv)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :])
+    # lora_output[:, output_dim_q: output_dim_q + output_dim_kv]
+    #      = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :])
+    # lora_output[:, output_dim_q + output_dim_kv: ]
+    #      = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = qkv_lora_b.shape[-1]
+    output_dim = qkv_lora_b.shape[-2]
+    assert input_dim == n_slices * r
+    assert output_offset.shape[0] == n_slices + 1
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S)
+        * triton.cdiv(max_qkv_out_dim, BLOCK_OUT),
+        n_slices,
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, output_dim), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    _qkv_lora_b_kernel[grid_b](
+        x,
+        qkv_lora_b,
+        output,
+        r,
+        max_qkv_out_dim,
+        x.stride(0),
+        x.stride(1),
+        qkv_lora_b.stride(0),
+        qkv_lora_b.stride(1),
+        qkv_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        output_offset,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/sgemm_lora_a.py

@@ -0,0 +1,182 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _sgemm_lora_a_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # stack_num * r
+    K,  # input_dim
+    stack_num,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA A matrix.
+
+    The kernel ensures that output[seg_start:seg_start + seg_len, :rank * stack_num]
+    stores the product of the input `x` and the LoRA weights for the corresponding
+    sequence. This implies that when rank is 0, the kernel is essentially a no-op,
+    as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty).
+
+    Args:
+        x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s`
+            is the sum of all sequence lengths in the batch.
+        weights (torch.Tensor): The LoRA 'A' weights for all available adapters,
+            with shape `(num_lora, N, K)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`.
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    batch_id = tl.program_id(axis=1)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel becomes a no-op as the output is always trivially correct.
+    if rank == 0:
+        return
+
+    pid = tl.program_id(axis=0)
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_len = tl.load(seg_lens + batch_id)
+    if seg_len == 0:
+        return
+
+    # Adjust N (stack_num * max_rank) according to the specific LoRA adapter
+    N = tl.minimum(N, rank * stack_num)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    # Create pointers for the first block of x and weights[batch_id]
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = x + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    w_ptrs = (weights + w_index * w_stride_0) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < N),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < N)
+    output_ptr = output + (
+        s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+    )
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def sgemm_lora_a_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    stack_num: int = 1,
+) -> torch.Tensor:
+    # x: (s, input_dim)
+    # weights: (num_lora, stack_num * r, input_dim)
+    # output: (s, stack_num * r)
+    # stack_num: run_qkv_lora: 3, run_gate_up_lora: 2
+    # when called by run_qkv_lora, the weights.shape[-2] will be 3 * r
+    # input_dim is much larger than r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    S = x.shape[0]
+    R = weights.shape[-2]
+    K = weights.shape[-1]
+    assert x.shape[-1] == K
+
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_K = 256
+    BLOCK_R = 16
+
+    grid = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(R, BLOCK_R),
+        batch_info.bs,
+    )
+
+    sorted_by_adapter = batch_info.permutation is not None
+
+    output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+    _sgemm_lora_a_kernel[grid](
+        x,
+        weights,
+        output,
+        R,
+        K,
+        stack_num,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_R,
+        BLOCK_K,
+    )
+    return output

python/sglang/kernels/ops/gemm/sgemm_lora_b.py

@@ -0,0 +1,188 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _sgemm_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # output_dim
+    K,  # r
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling
+    scalings,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA B matrix
+    and adds the result to the output in-place.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (torch.Tensor): The intermediate tensor from the LoRA 'A' multiplication,
+            of shape `(s, K)`, where `s` is the total number of tokens.
+        weights (torch.Tensor): The LoRA 'B' weights for all available adapters,
+            with shape `(num_lora, N, K)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`. This can be
+            the base model's output for a fused add operation.
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    batch_id = tl.program_id(axis=1)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    pid = tl.program_id(axis=0)
+    seg_len = tl.load(seg_lens + batch_id)
+    if seg_len == 0:
+        return
+    seg_start = tl.load(seg_indptr + batch_id)
+    scaling = tl.load(scalings + w_index)
+    # Adjust K (rank) according to the specific LoRA adapter
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    # Create pointers for the first block of x and weights[batch_id]
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = x + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    w_ptrs = (weights + w_index * w_stride_0) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iterate to compute the block in output matrix
+    n_mask = n_offset[None, :] < N
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K) & n_mask,
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = output + (
+        s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset[:, None] < seg_len) & n_mask
+    partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def sgemm_lora_b_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+    # x: (s, max_r)
+    # weights: (num_lora, output_dim, max_r)
+    # output: (s, output_dim)
+    # output_dim is much larger than max_r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    S = x.shape[0]
+    N = weights.shape[-2]
+    R = weights.shape[-1]
+    assert x.shape[-1] == R
+
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_N = 256
+
+    grid = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(N, BLOCK_N),
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((S, N), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    _sgemm_lora_b_kernel[grid](
+        x,
+        weights,
+        output,
+        N,
+        R,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_N,
+        BLOCK_R,
+        batch_info.scalings,
+    )
+    return output

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)."""

python/sglang/kernels/ops/gemm/trtllm_lora_temp/gate_up_lora_b.py

@@ -0,0 +1,266 @@
+import torch
+import triton
+import triton.language as tl
+
+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
+# the Triton kernel is faster (crossover measured at output_dim=1536/GPU:
+# cuBLAS wins rank64 from S>=256 and rank16 only from S>=2048).
+_CUBLAS_MIN_S_RANK = 16384
+
+
+@triton.jit
+def _gate_up_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    output_dim,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling
+    scalings,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    """
+    This kernel packs 2 sgemms (gate/up) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (Tensor): The input tensor, which is the result of the LoRA A projection.
+            Shape: (s, 2 * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, 2 * output_dim, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, 2 * output_dim).
+    """
+    # output_dim >> K
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # gate_up_id decides which of gate or up (0: gate, 1: up)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    gate_up_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+    seg_len = tl.load(seg_lens + batch_id)
+    if seg_len == 0:
+        return
+    seg_start = tl.load(seg_indptr + batch_id)
+    n_start = gate_up_id * output_dim  # offset on output dim
+    scaling = tl.load(scalings + w_index)
+
+    # Adjust K (rank) according to the specific LoRA adapter
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(output_dim, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    # Create pointers for the first block of x and weights
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = (
+        x
+        + (gate_up_id * K) * x_stride_1
+        + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    )
+    w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K)
+            & (n_offset[None, :] < output_dim),
+            other=0.0,
+        )
+        partial_sum += tl.dot(
+            x_tile.to(w_tile.dtype), w_tile
+        )  # cast fused: split-K returns fp32, plain path bf16 (no-op)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = (
+        output
+        + n_start * output_stride_1
+        + (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1)
+    )
+    output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < output_dim)
+    partial_sum += tl.load(output_ptr, mask=output_mask)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def _gate_up_lora_b_cublas(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Single-adapter dense path: one cuBLAS addmm_ per gate/up slice.
+
+    The LoRA-A output is rank-packed (slice i at columns [i*rank, (i+1)*rank)),
+    matching the Triton kernel's K = min(K, rank) slice stride. Slices are
+    disjoint output regions, so in-place addmm_ writes never collide.
+    """
+    r = gate_up_lora_b.shape[-1]
+    if base_output is None:
+        base_output = torch.zeros(
+            (x.shape[0], 2 * output_dim), device=x.device, dtype=x.dtype
+        )
+    w = gate_up_lora_b[0]
+    x_scaled = x[:, : 2 * r] * batch_info.scalings[0]
+    for i in range(2):
+        lo, hi = i * output_dim, (i + 1) * output_dim
+        base_output[:, lo:hi].addmm_(x_scaled[:, i * r : (i + 1) * r], w[lo:hi, :r].t())
+    return base_output
+
+
+def gate_up_lora_b_fwd(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+
+    # x: (s, 2 * r)
+    # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+    # output: (s, 2 * output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :])
+    # lora_output[:, output_dim:]
+    #      = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = gate_up_lora_b.shape[-1]
+    assert input_dim == 2 * r
+
+    if (
+        (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_GATE_UP.get()
+        )
+        and s * r >= _CUBLAS_MIN_S_RANK
+        and gate_up_lora_b.shape[0] == 1
+    ):  # single-adapter fast path: only valid with one resident slot
+        return _gate_up_lora_b_cublas(
+            x, gate_up_lora_b, batch_info, output_dim, base_output
+        )
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT),
+        2,  # this dimension decides current block computes on gate or up proj
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, 2 * output_dim), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _gate_up_lora_b_kernel[grid_b](
+        x,
+        gate_up_lora_b,
+        output,
+        r,
+        output_dim,
+        x.stride(0),
+        x.stride(1),
+        gate_up_lora_b.stride(0),
+        gate_up_lora_b.stride(1),
+        gate_up_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py

@@ -0,0 +1,31 @@
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.utils import is_arch_support_pdl
+
+
+def get_pdl_launch_metadata() -> tuple[bool, dict]:
+    """Return (ENABLE_PDL constexpr value, extra launch kwargs) for LoRA kernels.
+
+    ``launch_pdl`` is NVIDIA-only Triton launch metadata; the HIP backend
+    rejects unknown kwargs, so it is only included when PDL is supported.
+    """
+    enable_pdl = is_arch_support_pdl()
+    return enable_pdl, ({"launch_pdl": True} if enable_pdl else {})
+
+
+@triton.jit
+def _resolve_token_positions(
+    sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER: tl.constexpr
+):
+    """Map logical segment offsets to physical token positions.
+
+    When SORTED_BY_ADAPTER is True, segments are grouped by adapter and
+    sorted_token_ids provides the indirection to the original token rows.
+    When False, tokens are already contiguous starting at seg_start.
+    """
+    if SORTED_BY_ADAPTER:
+        return tl.load(
+            sorted_token_ids + seg_start + s_offset, mask=s_offset < seg_len
+        ).to(tl.int64)
+    return (seg_start + s_offset).to(tl.int64)

python/sglang/kernels/ops/gemm/trtllm_lora_temp/kv_b_lora_absorbed.py

@@ -0,0 +1,965 @@
+"""Triton kernels for absorbed-MLA ``kv_b_proj`` LoRA correction.
+
+The absorbed-MLA path bypasses ``kv_b_proj.forward()`` and folds the K/V
+sides as plain BMMs ``q_nope @ w_kc`` and ``attn_output @ w_vc``.  When a
+LoRA adapter is active on ``kv_b_proj`` we add the LoRA delta to
+``q_nope_out`` / ``attn_bmm_output`` manually.
+
+Using the standard LoRA factored math we *never* materialize ``B @ A``:
+
+    q_correction = q_nope     @ B_kc @ A   * scaling          # K-side
+    v_correction = attn_output @ A.T @ B_vc.T * scaling       # V-side
+
+where ``A: (slot, rank, kv_lora_rank)`` is the LoRA-A of ``kv_b_proj``
+(shared across heads) and ``B: (slot, num_heads*(qk_nope+v_head_dim), rank)``
+is the LoRA-B; ``B_kc`` / ``B_vc`` are its K-half / V-half slices.
+
+Four kernels split the math along the factorization boundary, all using
+the SGMM idiom from ``sgemm_lora_a`` / ``qkv_lora_b`` and the segment-indptr
+routing used by ``chunked_sgmv_*``:
+
+  * ``step_a_q_fwd``: per-head per-slot SGMM, ``(S,H,qk_nope) -> (S,H,rank)``
+  * ``step_b_q_fwd``: shared-A per-slot SGMM, scaled+accumulated,
+    ``(S,H,rank) -> (S,H,kv_lora_rank)``
+  * ``step_a_v_fwd``: shared-A.T per-slot SGMM, ``(S,H,kv_lora_rank) -> (S,H,rank)``
+  * ``step_b_v_fwd``: per-head per-slot SGMM with V-half of B, transposed,
+    scaled+accumulated, ``(S,H,rank) -> (S,H,v_head_dim)``
+
+Grid axes for each kernel:
+  axis 0 : output tile in (S, N)         -- tile_id = pid_s * num_pid_n + pid_n
+  axis 1 : head_id                       -- per-head weight slice
+  axis 2 : batch_id (segment / request)  -- per-slot weight routing via weight_indices
+
+Per-segment routing: each program derives its segment length from
+``seg_indptr[segment_id + 1] - seg_indptr[segment_id]``, loads
+``weight_indices[segment_id]`` once, and uses that slot's slice of the LoRA
+weight stack.  When ``permutation`` is present, rows are routed through it,
+matching the csgmv backend's adapter-grouped chunks.  No Python loops over slots
+or heads.
+
+The math also stays in the input dtype (no fp32 round-trip) -- the
+contraction dim ``rank`` is small (typically 16-64), so bf16 accumulation
+over it is acceptable.  ``tl.dot`` itself uses fp32 accumulation internally.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+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
+
+# ---------------------------------------------------------------------------
+# Block sizes -- chosen per-kernel from the natural shape of each step.
+#
+# The factored math gives the four kernels these contraction (K) and output
+# (N) ranges (for Kimi-K2.5: rank=16-32, qk_nope=v_head_dim=128, kv_lora_rank=512):
+#
+#                              K (contraction)      N (output)
+#   step_a_q                   qk_nope (~128)       rank (~16-32)
+#   step_b_q                   rank (~16-32)        kv_lora_rank (~512)
+#   step_a_v                   kv_lora_rank (~512)  rank (~16-32)
+#   step_b_v                   rank (~16-32)        v_head_dim (~128)
+#
+# So the "step_a_*" kernels want a large BLOCK_K (to keep loop iters small)
+# and a small BLOCK_N (matched to rank to avoid wasted tile lanes), while
+# the "step_b_*" kernels are the inverse.  Kernels aren't autotuned -- the
+# decode-shape workload is too small to benefit and the sweep surface is
+# wide.
+# ---------------------------------------------------------------------------
+
+_BLOCK_S = 16
+
+
+def _num_segments(batch_info: LoRABatchInfo) -> int:
+    return batch_info.num_segments or batch_info.bs
+
+
+def _max_segment_len(batch_info: LoRABatchInfo) -> int:
+    if batch_info.max_len is not None:
+        return batch_info.max_len
+    if batch_info.seg_lens is not None:
+        return int(batch_info.seg_lens.max().item())
+    raise ValueError("LoRA batch_info must provide max_len or seg_lens.")
+
+
+def _segment_grid_size(batch_info: LoRABatchInfo, num_segments: int) -> int:
+    return (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+
+
+# ---------------------------------------------------------------------------
+# Kernel 1 -- Step A_q: per-head per-slot SGMM, reads K-half of B
+#
+#     q_lora_a[t, h, r] = sum_{i<qk_nope} q_nope[t, h, i] * B[slot, h*FULL_K + i, r]
+#
+# x      : (S, H, qk_nope)
+# w (B)  : (num_lora, H*FULL_K, rank)   -- FULL_K = qk_nope + v_head_dim
+# out    : (S, H, rank)                 -- fresh allocation, no accumulate
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_q_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    H_FULL_K,  # H * (qk_nope + v_head_dim), the row-stride landmark
+    K,  # qk_nope (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,
+    w_stride_k,
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # per-head row stride in B (qk_nope + v_head_dim)
+    SORTED_BY_ADAPTER: tl.constexpr,
+    K_DIV: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank (allows mixed-rank batches).
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    # Clamp masked-lane indices into the valid range so pointer arithmetic
+    # stays in-bounds even before the load mask drops the values.
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    head_row_base = (
+        head_id * FULL_K
+    )  # row offset for this head's K-half (i in [0, qk_nope))
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = cur_k if K_DIV else tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + i, r]: row dim of B carries i (= GEMM K),
+        # column dim carries r (= GEMM N).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + (head_row_base + safe_k[:, None]) * w_stride_n
+            + safe_n[None, :] * w_stride_k,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_q_fwd(
+    q_nope: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    full_K_per_head: int,
+) -> torch.Tensor:
+    """Step A of the q-side correction.
+
+    Args:
+        q_nope: ``(S, H, qk_nope)``, the absorbed-MLA q intermediate.
+        B_buf: ``(num_lora, H*full_K_per_head, rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        full_K_per_head: ``qk_nope + v_head_dim``, the row stride per head in B.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate, ready for step B_q.
+    """
+    S, H, qk_nope_dim = q_nope.shape
+    rank = B_buf.shape[-1]
+
+    if (
+        lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+        or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
+    ) and B_buf.shape[
+        0
+    ] == 1:  # single-adapter fast path: only valid with one resident slot
+        # (S,H,r) view of a (H,S,r)-contiguous bmm result; step_b_q's dense
+        # path flattens in (h,s) order, so the chain needs no copies.
+        w_kc = B_buf[0].view(H, full_K_per_head, -1)[:, :qk_nope_dim, :]
+        return torch.bmm(q_nope.transpose(0, 1), w_kc).transpose(0, 1)
+
+    block_n = triton.next_power_of_2(rank)  # output N == rank -> one tile
+    out = torch.empty((S, H, rank), device=q_nope.device, dtype=q_nope.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_A_Q_BLOCK_K = 128
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, block_n),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_a_q_kernel[grid](
+        q_nope,
+        B_buf,
+        out,
+        S,
+        H * full_K_per_head,
+        qk_nope_dim,
+        rank,
+        q_nope.stride(0),
+        q_nope.stride(1),
+        q_nope.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        K_DIV=(qk_nope_dim % _STEP_A_Q_BLOCK_K == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=block_n,
+        BLOCK_K=_STEP_A_Q_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 2 -- Step B_q: shared-A per-slot SGMM, scaled + accumulated
+#
+#     base[t, h, k] += sum_r x[t, h, r] * A[slot, r, k] * scaling
+#
+# x      : (S, H, rank)
+# w (A)  : (num_lora, rank, kv_lora_rank)
+# base   : (S, H, kv_lora_rank), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_q_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # kv_lora_rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_k,
+    w_stride_n,
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    N_DIV: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    # Truncate contraction K to this slot's rank.
+    K_eff = tl.minimum(K, cur_rank)
+
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = n_offset if N_DIV else tl.minimum(n_offset, N - 1)
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]  (k iterates over rank)
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, k, n]: read k along contraction, n along output.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    # Accumulate into base[s, h, n].
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    tl.atomic_add(base + base_offs, partial_sum, mask=out_mask, sem="relaxed")
+
+
+def step_b_q_fwd(
+    q_lora_a: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Step B of the q-side correction, accumulating into ``base_output``.
+
+    Args:
+        q_lora_a: ``(S, H, rank)`` from step A_q.
+        A_buf: ``(num_lora, rank, kv_lora_rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, kv_lora_rank)``, modified in-place
+            (the absorbed ``q_nope @ w_kc`` result).
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = q_lora_a.shape
+    kv_lora_rank = A_buf.shape[-1]
+
+    if (
+        lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+        or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
+    ) and A_buf.shape[
+        0
+    ] == 1:  # single-adapter fast path: only valid with one resident slot
+        # Flatten (S,H) in whichever order base_output's storage allows
+        # without a copy (the absorbed q path passes a transpose view of a
+        # (H,S,kv)-contiguous bmm result). x is small; reshape may copy it.
+        base2d = x2d = None
+        if base_output.is_contiguous():
+            base2d = base_output.view(-1, kv_lora_rank)
+            x2d = q_lora_a[..., :rank].reshape(-1, rank)
+        elif base_output.transpose(0, 1).is_contiguous():
+            base2d = base_output.transpose(0, 1).view(-1, kv_lora_rank)
+            x2d = q_lora_a[..., :rank].transpose(0, 1).reshape(-1, rank)
+        if base2d is not None:
+            base2d.addmm_(x2d, A_buf[0, :, :], alpha=batch_info.scalings[0])
+            return base_output
+
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_B_Q_BLOCK_N = 128
+    _STEP_B_BLOCK_K = 16
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(kv_lora_rank, _STEP_B_Q_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_b_q_kernel[grid](
+        q_lora_a,
+        A_buf,
+        base_output,
+        S,
+        rank,
+        kv_lora_rank,
+        q_lora_a.stride(0),
+        q_lora_a.stride(1),
+        q_lora_a.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        N_DIV=(kv_lora_rank % _STEP_B_Q_BLOCK_N == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_Q_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return base_output
+
+
+# ---------------------------------------------------------------------------
+# Kernel 3 -- Step A_v: shared-A.T per-slot SGMM (no scaling, fresh output)
+#
+#     attn_lora_a[t, h, r] = sum_k attn_output[t, h, k] * A[slot, r, k]
+#
+# x      : (S, H, kv_lora_rank)
+# w (A)  : (num_lora, rank, kv_lora_rank) -- accessed transposed vs step B_q
+# out    : (S, H, rank), fresh allocation
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_v_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    K,  # kv_lora_rank (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # A's "rank" axis (= GEMM N)
+    w_stride_k,  # A's "kv_lora_rank" axis (= GEMM K)
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    K_DIV: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank.
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = cur_k if K_DIV else tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, r, k] -- here we want each (k, r) so we read along k
+        # (inner / contraction) and produce r as output.  Stride access:
+        # the row dim is r (= GEMM N), column dim is k (= GEMM K).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_v_fwd(
+    attn_output: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+) -> torch.Tensor:
+    """Step A of the v-side correction.
+
+    Args:
+        attn_output: ``(S, H, kv_lora_rank)``, the post-attention intermediate.
+        A_buf: ``(num_lora, rank, kv_lora_rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate for step B_v.
+    """
+    S, H, kv_lora_rank = attn_output.shape
+    rank = A_buf.shape[1]
+
+    if (
+        (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
+        )
+        and attn_output.is_contiguous()
+        and A_buf.shape[0] == 1
+    ):  # single-adapter fast path: only valid with one resident slot
+        return torch.mm(
+            attn_output.view(-1, kv_lora_rank), A_buf[0, :rank, :].t()
+        ).view(S, H, rank)
+
+    block_n = triton.next_power_of_2(rank)
+    out = torch.empty((S, H, rank), device=attn_output.device, dtype=attn_output.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_A_V_BLOCK_K = 256
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, block_n),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_a_v_kernel[grid](
+        attn_output,
+        A_buf,
+        out,
+        S,
+        kv_lora_rank,
+        rank,
+        attn_output.stride(0),
+        attn_output.stride(1),
+        attn_output.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        K_DIV=(kv_lora_rank % _STEP_A_V_BLOCK_K == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=block_n,
+        BLOCK_K=_STEP_A_V_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 4 -- Step B_v: per-head per-slot SGMM with V-half of B (transposed),
+# scaled + accumulated
+#
+#     base[t, h, j] += sum_r x[t, h, r] * B[slot, h*FULL_K + qk_nope + j, r] * scaling
+#
+# x      : (S, H, rank)
+# w (B)  : (num_lora, H*FULL_K, rank), V-half slice via offset
+# base   : (S, H, v_head_dim), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_v_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # v_head_dim (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # B's row dim (h*FULL_K + j) -- this is GEMM N
+    w_stride_k,  # B's rank dim -- this is GEMM K
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # qk_nope + v_head_dim
+    QK_NOPE_OFFSET: tl.constexpr,  # offset of V-half within each head's row block
+    SORTED_BY_ADAPTER: tl.constexpr,
+    N_DIV: tl.constexpr,  # N % BLOCK_N == 0 -> drop safe_n (keep the store coalesced)
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    K_eff = tl.minimum(K, cur_rank)
+
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = n_offset if N_DIV else tl.minimum(n_offset, N - 1)
+
+    # V-half row base for this head: h*FULL_K + qk_nope
+    head_row_base = head_id * FULL_K + QK_NOPE_OFFSET
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + qk_nope + j, r] -- row dim is j (= GEMM N),
+        # column dim is r (= GEMM K).  Transposed access vs step A_q.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + (head_row_base + safe_n[None, :]) * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    tl.atomic_add(base + base_offs, partial_sum, mask=out_mask, sem="relaxed")
+
+
+def step_b_v_fwd(
+    attn_lora_a: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+    qk_nope_head_dim: int,
+    v_head_dim: int,
+) -> torch.Tensor:
+    """Step B of the v-side correction, accumulating into ``base_output``.
+
+    Args:
+        attn_lora_a: ``(S, H, rank)`` from step A_v.
+        B_buf: ``(num_lora, H*(qk_nope+v_head_dim), rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, v_head_dim)``, modified in-place
+            (the absorbed ``attn_output @ w_vc`` result).
+        qk_nope_head_dim: offset of V-half within each head's row block of B.
+        v_head_dim: output feature dim per head.
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = attn_lora_a.shape
+    full_K_per_head = qk_nope_head_dim + v_head_dim
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_B_V_BLOCK_N = 64
+    _STEP_B_BLOCK_K = 16
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(v_head_dim, _STEP_B_V_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_b_v_kernel[grid](
+        attn_lora_a,
+        B_buf,
+        base_output,
+        S,
+        rank,
+        v_head_dim,
+        attn_lora_a.stride(0),
+        attn_lora_a.stride(1),
+        attn_lora_a.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        QK_NOPE_OFFSET=qk_nope_head_dim,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        N_DIV=(v_head_dim % _STEP_B_V_BLOCK_N == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_V_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return base_output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/qkv_lora_b.py

@@ -0,0 +1,297 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+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
+# this the Triton kernel is faster (measured crossover at the smallest
+# realistic N=768, where cuBLAS is weakest).
+_CUBLAS_MIN_MAX_LEN = 8
+
+
+@triton.jit
+def _qkv_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    max_qkv_out_dim,  # max(output_q_dim, output_kv_dim)
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    # Offsets of q/k/v slice on output dimension
+    n_offs,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling
+    scalings,
+    ENABLE_PDL: tl.constexpr = False,
+    STORE_WRITEBACK: tl.constexpr = False,
+):
+    """
+    This kernel packs 3 sgemms (q/k/v) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (Tensor): The input tensor, which is the result of the LoRA A projection.
+            Shape: (s, 3 * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank. The second dimension is partitioned
+            for Q, K, and V.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, N_Q + 2 * N_KV, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, N_Q + 2 * N_KV).
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    qkv_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+    seg_len = tl.load(seg_lens + batch_id)
+    if seg_len == 0:
+        return
+    seg_start = tl.load(seg_indptr + batch_id)
+    n_start = tl.load(n_offs + qkv_id)
+    n_size = tl.load(n_offs + qkv_id + 1) - n_start
+    scaling = tl.load(scalings + w_index)
+    # Adjust K (rank) according to the specific LoRA adapter.
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(max_qkv_out_dim, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    # Create pointers for the first block of x and weights[batch_id][n_start: n_end][:]
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = (
+        x
+        + (qkv_id * K) * x_stride_1
+        + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    )
+    w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    n_mask = n_offset[None, :] < n_size
+    x_tile = tl.load(
+        x_ptrs,
+        mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K),
+        other=0.0,
+    )
+    w_tile = tl.load(
+        w_ptrs,
+        mask=(k_offset[:, None] < K) & n_mask,
+        other=0.0,
+    )
+    # cast fused: the split-K shrink returns fp32, plain path bf16 (no-op)
+    partial_sum = tl.dot(x_tile.to(w_tile.dtype), w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix (cast to the OUTPUT dtype: x may be the fp32
+    # split-K shrink accumulator while base_output is bf16)
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(output.dtype.element_ty)
+    output_ptr = (
+        output
+        + n_start * output_stride_1
+        + (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1)
+    )
+    output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < n_size)
+    if STORE_WRITEBACK:
+        # The expand-add output tiles are disjoint across all programs in this launch
+        # (distinct s-rows / n-cols / slice / segment), so each element is RMW'd by
+        # exactly one program -- a plain read-add-write is correct and avoids the bf16
+        # narrow-tile atomic (the dominant decode cost). base_output is a same-stream
+        # data dependency (base GEMM before apply_lora), not a concurrent writer.
+        partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
+        tl.store(output_ptr, partial_sum, mask=output_mask)
+    else:
+        tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
+
+
+def _qkv_lora_b_cublas(
+    x: torch.Tensor,
+    qkv_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_offset_cpu: torch.Tensor,
+    base_output: Optional[torch.Tensor],
+    n_slices: int,
+) -> torch.Tensor:
+    """Single-adapter dense path: one cuBLAS addmm_ per q/k/v slice.
+
+    The LoRA-A output is rank-packed (slice i at columns [i*rank, (i+1)*rank)),
+    matching the Triton kernel's K = min(K, rank) slice stride. Slice offsets
+    come from the pinned CPU copy (no GPU sync); slices are disjoint output
+    regions, so in-place addmm_ writes never collide.
+    """
+    r = qkv_lora_b.shape[-1]
+    if base_output is None:
+        base_output = torch.zeros(
+            (x.shape[0], qkv_lora_b.shape[-2]), device=x.device, dtype=x.dtype
+        )
+    w = qkv_lora_b[0]
+    x_scaled = x[:, : n_slices * r] * batch_info.scalings[0]
+    offsets = output_offset_cpu.tolist()
+    for i in range(n_slices):
+        lo, hi = offsets[i], offsets[i + 1]
+        base_output[:, lo:hi].addmm_(x_scaled[:, i * r : (i + 1) * r], w[lo:hi, :r].t())
+    return base_output
+
+
+def qkv_lora_b_fwd(
+    x: torch.Tensor,
+    qkv_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_offset: torch.Tensor,
+    max_qkv_out_dim: int,
+    base_output: torch.Tensor = None,
+    n_slices: int = 3,
+    output_offset_cpu: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+
+    # x: (s, n_slices * r)
+    # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
+    # output_offset = [0, output_dim_q, output_dim_q + output_dim_kv,
+    #                     output_dim_q + 2 * output_dim_kv]  (length n_slices + 1)
+    # max_qkv_out_dim = max(output_dim_q, output_dim_kv)
+    # output: (s, output_dim_q + 2 * output_dim_kv)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :])
+    # lora_output[:, output_dim_q: output_dim_q + output_dim_kv]
+    #      = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :])
+    # lora_output[:, output_dim_q + output_dim_kv: ]
+    #      = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = qkv_lora_b.shape[-1]
+    output_dim = qkv_lora_b.shape[-2]
+    assert input_dim == n_slices * r
+    assert output_offset.shape[0] == n_slices + 1
+
+    if (
+        output_offset_cpu is not None
+        and (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_QKV.get()
+        )
+        and batch_info.max_len >= _CUBLAS_MIN_MAX_LEN
+        and qkv_lora_b.shape[0]
+        == 1  # single-adapter fast path: only valid with one resident slot
+    ):
+        return _qkv_lora_b_cublas(
+            x, qkv_lora_b, batch_info, output_offset_cpu, base_output, n_slices
+        )
+
+    BLOCK_S = 16
+    BLOCK_R = triton.next_power_of_2(r)
+    # BLOCK_OUT stays 64: with the 1-adapter cuBLAS dispatch the Triton path
+    # only runs for decode-sized batches, where 128 halves the grid (96->48
+    # programs on Kimi r16 bs64) and slows the kernel ~60% (11.4->18.5us, B200).
+    # Re-swept for the store path on GB200: 32 vs 64 is within noise (one preset
+    # marginally each way), so the single value is kept for both writebacks.
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S)
+        * triton.cdiv(max_qkv_out_dim, BLOCK_OUT),
+        n_slices,
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, output_dim), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    store_writeback = lora_envs.SGLANG_OPT_LORA_QKV_B_STORE.get()
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _qkv_lora_b_kernel[grid_b](
+        x,
+        qkv_lora_b,
+        output,
+        r,
+        max_qkv_out_dim,
+        x.stride(0),
+        x.stride(1),
+        qkv_lora_b.stride(0),
+        qkv_lora_b.stride(1),
+        qkv_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        output_offset,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+        ENABLE_PDL=enable_pdl,
+        STORE_WRITEBACK=store_writeback,
+        **pdl_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_a.py

@@ -0,0 +1,269 @@
+import functools
+
+import torch
+import triton
+import triton.language as tl
+
+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
+
+
+@triton.jit
+def _sgemm_lora_a_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # stack_num * r
+    K,  # input_dim
+    stack_num,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    SPLIT_K: tl.constexpr = 1,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA A matrix.
+
+    The kernel ensures that output[seg_start:seg_start + seg_len, :rank * stack_num]
+    stores the product of the input `x` and the LoRA weights for the corresponding
+    sequence. This implies that when rank is 0, the kernel is essentially a no-op,
+    as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty).
+
+    Args:
+        x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s`
+            is the sum of all sequence lengths in the batch.
+        weights (torch.Tensor): The LoRA 'A' weights for all available adapters,
+            with shape `(num_lora, N, K)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`.
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    batch_id = tl.program_id(axis=1)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel becomes a no-op as the output is always trivially correct.
+    if rank == 0:
+        return
+
+    pid = tl.program_id(axis=0)
+    # Fold the split-K factor out of axis-0 (SPLIT_K == 1 -> pid_sk == 0, pid_tile == pid).
+    pid_sk = pid % SPLIT_K
+    pid_tile = pid // SPLIT_K
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_len = tl.load(seg_lens + batch_id)
+    if seg_len == 0:
+        return
+
+    # Adjust N (stack_num * max_rank) to this adapter's actual rank.
+    N = tl.minimum(N, rank * stack_num)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid_tile // num_pid_n
+    pid_n = pid_tile % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    # Create pointers for the first block of x and weights[batch_id]
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate.
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = x + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    w_ptrs = (weights + w_index * w_stride_0) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
+        k_remaining = K - k * (BLOCK_K * SPLIT_K)
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < k_remaining),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < k_remaining) & (n_offset[None, :] < N),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * SPLIT_K * x_stride_1
+        w_ptrs += BLOCK_K * SPLIT_K * w_stride_2
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix
+    output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < N)
+    output_ptr = output + (
+        s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+    )
+    if SPLIT_K == 1:
+        tl.store(output_ptr, partial_sum.to(output.dtype.element_ty), mask=output_mask)
+    else:
+        tl.atomic_add(
+            output_ptr,
+            partial_sum.to(output.dtype.element_ty),
+            mask=output_mask,
+            sem="relaxed",
+        )
+
+
+@functools.lru_cache(maxsize=None)
+def _num_sms(device_index: int) -> int:
+    return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def sgemm_lora_a_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    stack_num: int = 1,
+    out_alloc_stream=None,
+) -> torch.Tensor:
+    # x: (s, input_dim)
+    # weights: (num_lora, stack_num * r, input_dim)
+    # output: (s, stack_num * r)
+    # stack_num: run_qkv_lora: 3, run_gate_up_lora: 2
+    # when called by run_qkv_lora, the weights.shape[-2] will be 3 * r
+    # input_dim is much larger than r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    S = x.shape[0]
+    R = weights.shape[-2]
+    K = weights.shape[-1]
+    assert x.shape[-1] == K
+
+    if (
+        lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+        or lora_envs.SGLANG_OPT_LORA_CUBLAS_A.get()
+    ) and weights.shape[
+        0
+    ] == 1:  # single-adapter fast path: only valid with one resident slot
+        # Honor out_alloc_stream like the Triton path below: under SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC
+        # the shrink output must be allocated on the MAIN (consumer) stream so the caching allocator
+        # frees/reuses it on the consumer's schedule (cuda-graph WAR). F.linear has no out=, so
+        # allocate explicitly and matmul into it.
+        if out_alloc_stream is not None:
+            with torch.cuda.stream(out_alloc_stream):
+                output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+        else:
+            output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+        return torch.matmul(x, weights[0].transpose(-2, -1), out=output)
+
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_K = 256
+    BLOCK_R = triton.next_power_of_2(R)
+
+    sorted_by_adapter = batch_info.permutation is not None
+
+    num_s_tiles = triton.cdiv(batch_info.max_len, BLOCK_S)
+    split_k = 1
+    if lora_envs.SGLANG_ENABLE_LORA_SHRINK_SPLIT_K.get() and x.is_cuda:
+        num_k_tiles = triton.cdiv(K, BLOCK_K)
+        base_grid = batch_info.bs * num_s_tiles
+        num_sms = _num_sms(x.device.index)
+        if base_grid < num_sms and num_k_tiles >= 16:
+            split_k = max(1, min(2 * num_sms // base_grid, num_k_tiles, 16))
+
+    launch_kwargs = {}
+    if split_k > 1:
+        # out_alloc_stream (SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC) is intentionally NOT honored here:
+        # torch.zeros launches its memset on the alloc stream, which would race the side-stream
+        # shrink without extra ordering. No current config exercises split-K together with the
+        # two-stream main-alloc overlap (qwen3.5 leaves split-K off; kimi is single-stream-coherent).
+        output = torch.zeros((S, R), device=x.device, dtype=torch.float32)
+        launch_kwargs = {
+            "num_warps": 2 if split_k <= 4 else 4,
+            "num_stages": 3,
+        }
+    elif out_alloc_stream is not None:
+        # Allocate the output on the MAIN (consumer) stream when requested, so the caching allocator
+        # frees/reuses it on the consumer's schedule, not the side stream's (cuda-graph WAR — see
+        # lora_overlap_alloc_stream / SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC).
+        with torch.cuda.stream(out_alloc_stream):
+            output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+    else:
+        output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+
+    grid = (
+        num_s_tiles * split_k,
+        batch_info.bs,
+    )
+
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _sgemm_lora_a_kernel[grid](
+        x,
+        weights,
+        output,
+        R,
+        K,
+        stack_num,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_R,
+        BLOCK_K,
+        split_k,
+        ENABLE_PDL=enable_pdl,
+        **launch_kwargs,
+        **pdl_kwargs,
+    )
+    # split_k>1 returns the fp32 accumulator directly; the LoRA-B expand casts x to the weight dtype
+    # on-load (fused), dropping the standalone fp32->bf16 copy kernel. split_k==1 already returns x.dtype.
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_b.py

@@ -0,0 +1,223 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.trtllm_lora_temp.gate_up_lora_b import (
+    _CUBLAS_MIN_S_RANK,
+)
+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
+
+
+def _sgemm_lora_b_cublas(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Single-adapter dense path: one cuBLAS addmm_ over the full output.
+
+    Mirrors the Triton kernel exactly: single slice, K = max_r (the kernel
+    reads the full K; the loader zero-pads the weight tail beyond the
+    adapter's rank), scaling fused via a pre-scaled x.
+    """
+    if base_output is None:
+        base_output = torch.zeros(
+            (x.shape[0], weights.shape[-2]), device=x.device, dtype=x.dtype
+        )
+    w = weights[batch_info.weight_indices[0]]
+    x_scaled = x * batch_info.scalings[0]
+    base_output.addmm_(x_scaled, w.t())
+    return base_output
+
+
+@triton.jit
+def _sgemm_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # output_dim
+    K,  # r
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    # Meta parameters
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling
+    scalings,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA B matrix
+    and adds the result to the output in-place.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (torch.Tensor): The intermediate tensor from the LoRA 'A' multiplication,
+            of shape `(s, K)`, where `s` is the total number of tokens.
+        weights (torch.Tensor): The LoRA 'B' weights for all available adapters,
+            with shape `(num_lora, N, K)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`. This can be
+            the base model's output for a fused add operation.
+    """
+
+    pid_s = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    seg_len = tl.load(seg_lens + batch_id)
+    if pid_s * BLOCK_S >= seg_len:  # also covers seg_len == 0
+        return
+    seg_start = tl.load(seg_indptr + batch_id)
+    scaling = tl.load(scalings + w_index)
+
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+    x_ptrs = x + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    w_ptrs = (weights + w_index * w_stride_0) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    n_mask = n_offset[None, :] < N
+    output_ptr = output + (
+        s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset[:, None] < seg_len) & n_mask
+
+    x_tile = tl.load(
+        x_ptrs,
+        mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K),
+        other=0.0,
+    )
+    w_tile = tl.load(
+        w_ptrs,
+        mask=(k_offset[:, None] < K) & n_mask,
+        other=0.0,
+    )
+
+    # cast fused: the split-K shrink returns fp32, plain path bf16 (no-op)
+    partial_sum = tl.dot(x_tile.to(w_tile.dtype), w_tile) * scaling
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix (cast to the OUTPUT dtype: x may be the fp32
+    # split-K shrink accumulator while base_output is bf16)
+    partial_sum = partial_sum.to(output.dtype.element_ty)
+    tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
+
+
+def sgemm_lora_b_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+    # x: (s, max_r)
+    # weights: (num_lora, output_dim, max_r)
+    # output: (s, output_dim)
+    # output_dim is much larger than max_r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    S = x.shape[0]
+    N = weights.shape[-2]
+    R = weights.shape[-1]
+    assert x.shape[-1] == R
+
+    if (
+        (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_B.get()
+        )
+        and S * R >= _CUBLAS_MIN_S_RANK
+        and weights.shape[0] == 1
+    ):  # single-adapter fast path: only valid with one resident slot
+        return _sgemm_lora_b_cublas(x, weights, batch_info, base_output)
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_R = triton.next_power_of_2(R)
+    BLOCK_N = 256
+
+    grid = (
+        triton.cdiv(batch_info.max_len, BLOCK_S),
+        triton.cdiv(N, BLOCK_N),
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((S, N), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _sgemm_lora_b_kernel[grid](
+        x,
+        weights,
+        output,
+        N,
+        R,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_N,
+        BLOCK_R,
+        batch_info.scalings,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return output

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.<module>``); 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__ = []

python/sglang/kernels/ops/grammar/bitmask_ops.py

@@ -0,0 +1,141 @@
+# Adapt from
+# https://github.com/mlc-ai/xgrammar/blob/v0.1.17/python/xgrammar/kernels/apply_token_bitmask_inplace_triton.py
+
+from typing import List, Optional, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import get_device_core_count
+
+
+@triton.jit
+def apply_token_bitmask_inplace_kernel(
+    logits_ptr,
+    bitmask_ptr,
+    indices_ptr,
+    num_rows,
+    vocab_size,
+    logits_strides,
+    bitmask_strides,
+    NUM_SMS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Apply a bitmask to logits in-place using Triton. The bitmask is a 01 bitwise compressed tensor,
+    where 0 means the token is masked and 1 means the token is not masked. After applying the bitmask,
+    the masked logits will be set to -inf.
+
+    Parameters
+    ----------
+    logits_ptr : tl.tensor
+        Pointer to the logits tensor to apply the bitmask to.
+
+    bitmask_ptr : tl.tensor
+        Pointer to the bitmask tensor to apply.
+
+    indices_ptr : Optional[tl.tensor]
+        Optional pointer to indices tensor specifying which rows to apply the mask to.
+
+    num_rows : int
+        Number of rows to process. If indices_ptr is provided, this is the number of unique indices.
+
+    vocab_size : int
+        Size of the vocabulary dimension. If the logits does not have a vocab padding, this is the
+        same as the logits's second dimension. Otherwise, this is the actual size of the vocabulary.
+
+    logits_strides : int
+        Stride between rows in the logits tensor.
+
+    bitmask_strides : int
+        Stride between rows in the bitmask tensor.
+
+    NUM_SMS : int
+        Number of streaming multiprocessors to use.
+
+    BLOCK_SIZE : int
+        Size of processing blocks.
+    """
+
+    pid = tl.program_id(0)
+    num_blocks = tl.cdiv(vocab_size, BLOCK_SIZE)
+    for work_id in tl.range(pid, num_rows * num_blocks, NUM_SMS):
+        row_id = work_id // num_blocks
+        block_offset = (work_id % num_blocks) * BLOCK_SIZE
+        batch_id = row_id if indices_ptr is None else tl.load(indices_ptr + row_id)
+        offsets = block_offset + tl.arange(0, BLOCK_SIZE)
+        bitmask_offsets = block_offset // 32 + tl.arange(0, BLOCK_SIZE // 32)
+        vocab_mask = offsets < vocab_size
+        packed_bitmask_mask = bitmask_offsets < bitmask_strides
+        packed_bitmask = tl.load(
+            bitmask_ptr + batch_id * bitmask_strides + bitmask_offsets,
+            packed_bitmask_mask,
+        )
+        bitmask = ((packed_bitmask[:, None] >> (tl.arange(0, 32)[None, :])) & 1) == 0
+        bitmask = bitmask.reshape(BLOCK_SIZE)
+
+        tl.store(
+            logits_ptr + batch_id * logits_strides + offsets,
+            -float("inf"),
+            vocab_mask & bitmask,
+        )
+
+
+def apply_token_bitmask_inplace_triton(
+    logits: torch.Tensor,
+    bitmask: torch.Tensor,
+    indices: Optional[Union[List[int], torch.Tensor]] = None,
+):
+    NUM_SMS = get_device_core_count()
+    BLOCK_SIZE = 4096
+    BITS_PER_BLOCK = 32
+
+    # Check input dtype
+    assert bitmask.dtype == torch.int32, "bitmask must be of type int32"
+
+    # Check input tensor shapes.
+    logits_shape = logits.shape
+    bitmask_shape = bitmask.shape
+    if logits.ndim == 1:
+        logits_shape = (1, logits_shape[0])
+    if bitmask.ndim == 1:
+        bitmask_shape = (1, bitmask_shape[0])
+
+    required_bitmask_width = (logits_shape[1] + BITS_PER_BLOCK - 1) // BITS_PER_BLOCK
+    assert required_bitmask_width >= bitmask_shape[1], (
+        f"Bitmask width too large: allow at most {required_bitmask_width} int32s for "
+        f"logits' width {logits_shape[1]}, but got {bitmask_shape[1]}"
+    )
+
+    vocab_size = min(logits_shape[1], bitmask_shape[1] * BITS_PER_BLOCK)
+
+    num_rows = None
+    if isinstance(indices, list) or isinstance(indices, torch.Tensor):
+        indices = torch.tensor(indices, dtype=torch.int32, device=logits.device)
+        num_rows = indices.shape[0]
+    else:
+        assert (
+            logits_shape[0] == bitmask_shape[0]
+        ), f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}"
+        num_rows = logits_shape[0]
+
+    if NUM_SMS > 0:
+        grid = (NUM_SMS,)
+    else:
+        num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
+        grid = (num_rows * num_blocks,)
+        NUM_SMS = triton.next_power_of_2(grid[0])
+
+    apply_token_bitmask_inplace_kernel[grid](
+        logits,
+        bitmask,
+        indices,
+        num_rows,
+        vocab_size,
+        logits_shape[1],
+        bitmask_shape[1],
+        NUM_SMS,
+        BLOCK_SIZE,
+        num_warps=BLOCK_SIZE // 32 // (16 // logits.element_size()),
+        num_stages=3,
+    )

python/sglang/kernels/ops/grammar/token_filter_ops.py

@@ -0,0 +1,175 @@
+# Copyright 2026 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 token filter operations."""
+
+from collections import OrderedDict
+from typing import List
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import get_device_core_count
+
+
+@triton.jit
+def reset_vocab_mask_kernel(
+    vocab_mask_ptr,
+    batch_idx: int,
+    num_elements: int,
+    reset_value: tl.constexpr,
+):
+    """Reset the vocab mask for a specific batch index to a given value.
+
+    Parameters
+    ----------
+    vocab_mask_ptr : tl.tensor
+        Pointer to the vocab mask tensor.
+
+    batch_idx : int
+        The batch index to reset.
+
+    num_elements : int
+        Number of int32 elements in the vocab mask for each batch.
+
+    reset_value : int
+        The value to reset the vocab mask to (typically -1 or 0).
+    """
+    pid = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for i in tl.range(pid, num_elements, num_threads):
+        offset = batch_idx * num_elements + i
+        tl.store(vocab_mask_ptr + offset, reset_value)
+
+
+@triton.jit
+def set_token_filter_batch_kernel(
+    vocab_mask_ptr,
+    token_ids_ptr,
+    batch_idx: int,
+    num_tokens: int,
+    num_elements: int,
+    is_allowed: tl.constexpr,
+):
+    """Set or clear specific tokens in the vocab mask for a batch.
+
+    Each token ID maps to a specific bit in the int32 bitmask array.
+    The kernel sets or clears those bits using atomic operations.
+
+    Parameters
+    ----------
+    vocab_mask_ptr : tl.tensor
+        Pointer to the vocab mask tensor.
+
+    token_ids_ptr : tl.tensor
+        Pointer to the token IDs to set/clear.
+
+    batch_idx : int
+        The batch index to modify.
+
+    num_tokens : int
+        Number of tokens to process.
+
+    num_elements : int
+        Number of int32 elements in the vocab mask for each batch.
+
+    is_allowed : bool
+        If True, set the bit to 1 (allow token).
+        If False, clear the bit to 0 (block token).
+    """
+    pid = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for i in tl.range(pid, num_tokens, num_threads):
+        token_id = tl.load(token_ids_ptr + i)
+        element_idx = token_id // 32
+        bit_idx = token_id % 32
+
+        offset = batch_idx * num_elements + element_idx
+
+        if is_allowed:
+            tl.atomic_or(vocab_mask_ptr + offset, 1 << bit_idx)
+        else:
+            tl.atomic_and(vocab_mask_ptr + offset, ~(1 << bit_idx))
+
+
+_cached_num_sms = None
+_cached_token_id_tensors: OrderedDict[tuple[int, tuple[int, ...]], torch.Tensor] = (
+    OrderedDict()
+)
+_MAX_TOKEN_ID_TENSOR_CACHE_SIZE = 32
+
+
+def _compute_grid(work_items: int):
+    global _cached_num_sms
+    if _cached_num_sms is None:
+        _cached_num_sms = get_device_core_count()
+    if _cached_num_sms > 0:
+        return (min(_cached_num_sms, work_items),)
+    return (work_items,)
+
+
+def _get_cached_token_ids_tensor(
+    token_ids: List[int], device: torch.device
+) -> torch.Tensor:
+    key = (device.index or 0, tuple(token_ids))
+    cached = _cached_token_id_tensors.get(key)
+    if cached is not None:
+        _cached_token_id_tensors.move_to_end(key)
+        return cached
+
+    token_ids_tensor = torch.tensor(token_ids, dtype=torch.int32, device=device)
+    _cached_token_id_tensors[key] = token_ids_tensor
+    if len(_cached_token_id_tensors) > _MAX_TOKEN_ID_TENSOR_CACHE_SIZE:
+        _cached_token_id_tensors.popitem(last=False)
+    return token_ids_tensor
+
+
+def set_token_filter_triton(
+    vocab_mask: torch.Tensor,
+    token_ids: List[int],
+    batch_idx: int,
+    is_allowed: bool = True,
+    reset_vocab_mask: bool = True,
+):
+    """Set or clear specific tokens in the vocab mask using Triton."""
+    assert vocab_mask.device.type == "cuda"
+
+    num_elements = vocab_mask.shape[1]
+
+    if reset_vocab_mask:
+        reset_value = 0 if is_allowed else -1
+        reset_vocab_mask_kernel[_compute_grid(num_elements)](
+            vocab_mask,
+            batch_idx,
+            num_elements,
+            reset_value,
+            num_warps=4,
+        )
+
+    if not token_ids:
+        return
+
+    num_tokens = len(token_ids)
+    token_ids_tensor = _get_cached_token_ids_tensor(token_ids, vocab_mask.device)
+    set_token_filter_batch_kernel[_compute_grid(num_tokens)](
+        vocab_mask,
+        token_ids_tensor,
+        batch_idx,
+        num_tokens,
+        num_elements,
+        is_allowed,
+        num_warps=4,
+    )

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

python/sglang/kernels/ops/kvcache/aiter_unified_attention.py

@@ -0,0 +1,97 @@
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def scatter_ragged_to_page_table_kernel(
+    kv_flat_ptr,
+    kv_indptr_ptr,
+    dest_ptr,
+    dest_stride,
+    sw_page_table_ptr,
+    swa_slot_mapping_ptr,
+    PAGE_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+):
+    """Scatter ragged token-level kv_indices into a 2D block-level page table."""
+    pid = tl.program_id(0)
+    block_id = tl.program_id(1)
+
+    start = tl.load(kv_indptr_ptr + pid).to(tl.int64)
+    kv_len = tl.load(kv_indptr_ptr + pid + 1).to(tl.int64) - start
+    num_blocks = (kv_len + PAGE_SIZE - 1) // PAGE_SIZE
+
+    offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    if block_id * BLOCK_SIZE >= num_blocks:
+        return
+    mask = offsets < num_blocks
+    token_idx = offsets.to(tl.int64) * PAGE_SIZE
+    vals = tl.load(kv_flat_ptr + start + token_idx, mask=mask, other=0)
+    block_vals = vals // PAGE_SIZE
+    tl.store(
+        dest_ptr + pid.to(tl.int64) * dest_stride + offsets,
+        block_vals,
+        mask=mask,
+    )
+
+    if HAS_SWA:
+        sw_vals = tl.load(swa_slot_mapping_ptr + vals)
+        block_vals = sw_vals // PAGE_SIZE
+        tl.store(
+            sw_page_table_ptr + pid.to(tl.int64) * dest_stride + offsets,
+            block_vals,
+            mask=mask,
+        )
+
+
+@triton.jit
+def scatter_req_to_token_to_page_table_kernel(
+    req_to_token_ptr,
+    req_pool_indices_ptr,
+    seq_lens_ptr,
+    page_table_ptr,
+    req_to_token_stride,
+    page_table_stride,
+    sw_page_table_ptr,
+    swa_slot_mapping_ptr,
+    DRAFT_NUM: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+):
+    """Build the 2D block-level page_table for target_verify from req_to_token."""
+    pid = tl.program_id(0)
+    block_id = tl.program_id(1)
+
+    seq_len = tl.load(seq_lens_ptr + pid).to(tl.int64)
+    kv_len = seq_len + DRAFT_NUM
+    num_blocks = (kv_len + PAGE_SIZE - 1) // PAGE_SIZE
+
+    offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    if block_id * BLOCK_SIZE >= num_blocks:
+        return
+    mask = offsets < num_blocks
+
+    rp = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+    token_idx = offsets.to(tl.int64) * PAGE_SIZE
+    vals = tl.load(
+        req_to_token_ptr + rp * req_to_token_stride + token_idx,
+        mask=mask,
+        other=0,
+    )
+    block_vals = vals // PAGE_SIZE
+    tl.store(
+        page_table_ptr + pid.to(tl.int64) * page_table_stride + offsets,
+        block_vals,
+        mask=mask,
+    )
+
+    if HAS_SWA:
+        sw_vals = tl.load(swa_slot_mapping_ptr + vals)
+        block_vals = sw_vals // PAGE_SIZE
+        tl.store(
+            sw_page_table_ptr + pid.to(tl.int64) * page_table_stride + offsets,
+            block_vals,
+            mask=mask,
+        )

python/sglang/kernels/ops/kvcache/cache_move.py

@@ -0,0 +1,306 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cpu
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import copy_all_layer_kv_cache_cpu
+
+
+@triton.jit
+def set_kv_buffer_prefix_valid_tiled(
+    src_k_ptr,
+    src_v_ptr,
+    dst_k_ptr,
+    dst_v_ptr,
+    loc_2d_ptr,
+    commit_len_ptr,
+    src_k_row_stride,
+    src_v_row_stride,
+    dst_k_row_stride,
+    dst_v_row_stride,
+    block_size,
+    ROW_BYTES: tl.constexpr,
+    BYTES_PER_TILE: tl.constexpr,
+):
+    bid = tl.program_id(0)
+    row = tl.program_id(1)
+    tid = tl.program_id(2)
+
+    commit_len = tl.load(commit_len_ptr + bid)
+    if row >= commit_len:
+        return
+
+    byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE)
+    mask_byte = byte_off < ROW_BYTES
+    tl.multiple_of(byte_off, 16)
+
+    loc = tl.load(loc_2d_ptr + bid * block_size + row)
+    src_row = bid * block_size + row
+
+    src_k_ptr = tl.cast(src_k_ptr, tl.pointer_type(tl.uint8))
+    src_v_ptr = tl.cast(src_v_ptr, tl.pointer_type(tl.uint8))
+    dst_k_ptr = tl.cast(dst_k_ptr, tl.pointer_type(tl.uint8))
+    dst_v_ptr = tl.cast(dst_v_ptr, tl.pointer_type(tl.uint8))
+
+    src_k_row_ptr = src_k_ptr + src_row * src_k_row_stride + byte_off
+    src_v_row_ptr = src_v_ptr + src_row * src_v_row_stride + byte_off
+    dst_k_row_ptr = dst_k_ptr + loc * dst_k_row_stride + byte_off
+    dst_v_row_ptr = dst_v_ptr + loc * dst_v_row_stride + byte_off
+
+    k_val = tl.load(src_k_row_ptr, mask=mask_byte, other=0)
+    v_val = tl.load(src_v_row_ptr, mask=mask_byte, other=0)
+    tl.store(dst_k_row_ptr, k_val, mask=mask_byte)
+    tl.store(dst_v_row_ptr, v_val, mask=mask_byte)
+
+
+@triton.jit
+def copy_all_layer_kv_cache_tiled(
+    data_ptrs,
+    strides,
+    tgt_loc_ptr,
+    src_loc_ptr,
+    num_locs,
+    num_locs_upper: tl.constexpr,
+    BYTES_PER_TILE: tl.constexpr,
+):
+    """2D tiled kernel. Safe for in-place copy."""
+    bid = tl.program_id(0)
+    tid = tl.program_id(1)
+
+    stride = tl.load(strides + bid)
+    base_ptr = tl.load(data_ptrs + bid)
+    base_ptr = tl.cast(base_ptr, tl.pointer_type(tl.uint8))
+
+    byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE)
+    mask_byte = byte_off < stride
+    tl.multiple_of(byte_off, 16)
+
+    loc_idx = tl.arange(0, num_locs_upper)
+    mask_loc = loc_idx < num_locs
+
+    src = tl.load(src_loc_ptr + loc_idx, mask=mask_loc, other=0)
+    tgt = tl.load(tgt_loc_ptr + loc_idx, mask=mask_loc, other=0)
+
+    src_ptr = base_ptr + src[:, None] * stride + byte_off[None, :]
+    tgt_ptr = base_ptr + tgt[:, None] * stride + byte_off[None, :]
+
+    mask = mask_loc[:, None] & mask_byte[None, :]
+    vals = tl.load(src_ptr, mask=mask)
+    tl.store(tgt_ptr, vals, mask=mask)
+
+
+def copy_all_layer_kv_cache_func(
+    data_ptrs: torch.Tensor,
+    strides: torch.Tensor,
+    tgt_loc: torch.Tensor,
+    src_loc: torch.Tensor,
+    num_locs: int,
+    num_locs_upper: int,
+    kv_copy_config: dict,
+):
+    if _is_cpu:
+        copy_all_layer_kv_cache_cpu(
+            data_ptrs,
+            strides,
+            tgt_loc[:num_locs],
+            src_loc[:num_locs],
+        )
+        return
+    grid = (data_ptrs.numel(), kv_copy_config["byte_tiles"])
+    copy_all_layer_kv_cache_tiled[grid](
+        data_ptrs,
+        strides,
+        tgt_loc,
+        src_loc,
+        num_locs,
+        num_locs_upper,
+        BYTES_PER_TILE=kv_copy_config["bytes_per_tile"],
+        num_warps=kv_copy_config["num_warps"],
+        num_stages=2,
+    )
+
+
+# ---------------------------------------------------------------------------
+# store_cache_4d — single-launch Triton write into the 4-D page-major envelope
+# K/V views. At `PAGE_SIZE = 1` the kernel constexpr-folds to byte-identical
+# addresses as the slot-major envelope view; at `PAGE_SIZE > 1` it uses the
+# same `(page_id, tok_in_p)` split the attention read kernels use.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def store_cache_4d_kernel(
+    k_view_ptr,
+    v_view_ptr,
+    cache_k_ptr,
+    cache_v_ptr,
+    loc_ptr,
+    # Strides in ELEMENTS (not bytes); wrapper passes view.stride(D)
+    # directly. K and V may have different head_dim → different per-token
+    # strides, so we carry both.
+    stride_k_page,
+    stride_k_tok,
+    stride_v_page,
+    stride_v_tok,
+    stride_src_k_row,
+    stride_src_v_row,
+    K_ROW_DIM: tl.constexpr,  # head_num * head_dim
+    V_ROW_DIM: tl.constexpr,  # head_num * v_head_dim
+    PAGE_SIZE: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Token-parallel Triton write into a 4-D envelope-strided K/V view.
+
+    Grid: ``(N, ceil(max(K_ROW_DIM, V_ROW_DIM) / BLOCK), 2)`` where:
+      - axis 0 → one program per token (loc[i])
+      - axis 1 → blocks within one slot's K (or V) row
+      - axis 2 → 0 = K, 1 = V (two-tensor write fused into one launch)
+
+    For each token i, the kernel writes:
+        page_id  = loc[i] // PAGE_SIZE
+        tok_in_p = loc[i] %  PAGE_SIZE
+        k_view[page_id, tok_in_p, :, :] = cache_k[i, :, :]
+        v_view[page_id, tok_in_p, :, :] = cache_v[i, :, :]
+
+    Cuda-graph safe: no Python branching on tensor values, no `.item()`,
+    all shapes/strides known at launch time.
+    """
+    pid_n = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_kv = tl.program_id(2)
+
+    # 1. Resolve destination slot in the 4-D view.
+    loc = tl.load(loc_ptr + pid_n).to(tl.int64)
+    if PAGE_SIZE == 1:
+        page_id = loc
+        tok_in_p = tl.zeros([], dtype=tl.int64)
+    else:
+        page_id = loc // PAGE_SIZE
+        tok_in_p = loc % PAGE_SIZE
+
+    # 2. Compute per-tensor source/dest pointers.
+    base_off = pid_b * BLOCK + tl.arange(0, BLOCK)
+
+    if pid_kv == 0:
+        mask = base_off < K_ROW_DIM
+        # The trailing (head_num, head_dim) axes of `k_view` are
+        # contiguous: stride[-1]==1, stride[-2]==head_dim. So we can
+        # treat them as a flat K_ROW_DIM dimension addressed by `base_off`.
+        # The wrapper asserts this invariant.
+        src_ptr = cache_k_ptr + pid_n * stride_src_k_row + base_off
+        dst_ptr = (
+            k_view_ptr + page_id * stride_k_page + tok_in_p * stride_k_tok + base_off
+        )
+    else:
+        mask = base_off < V_ROW_DIM
+        src_ptr = cache_v_ptr + pid_n * stride_src_v_row + base_off
+        dst_ptr = (
+            v_view_ptr + page_id * stride_v_page + tok_in_p * stride_v_tok + base_off
+        )
+
+    src = tl.load(src_ptr, mask=mask)
+    tl.store(dst_ptr, src, mask=mask)
+
+
+def store_cache_4d(
+    k_view: torch.Tensor,
+    v_view: torch.Tensor,
+    cache_k: torch.Tensor,
+    cache_v: torch.Tensor,
+    loc: torch.Tensor,
+    page_size: int,
+) -> None:
+    """One-launch Triton write into the 4-D page-major envelope K/V views.
+
+    Writes ``cache_k[i]`` and ``cache_v[i]`` to
+    ``k_view[loc[i]//ps, loc[i]%ps, :, :]`` (and analogously for V) for
+    ``i in [0, N)``.
+
+    Contract:
+        - ``k_view``, ``v_view``: 4-D ``(num_pages, page_size, head_num,
+          head_dim*)``, contiguous in the trailing ``(head_num, head_dim)``
+          dims (i.e., ``stride[-1] == 1`` and ``stride[-2] == head_dim``).
+        - ``cache_k``, ``cache_v``: 3-D ``(N, head_num, head_dim*)``,
+          contiguous in the trailing ``(head_num, head_dim)`` dims.
+        - ``loc``: 1-D int64 or int32, N elements, values in
+          ``[0, num_pages * page_size)``. The caller is responsible for
+          clamping any negative entries to ≥ 0.
+        - At ``page_size == 1`` the kernel produces byte-identical output
+          to the legacy advanced-indexing path.
+
+    Returns nothing; writes in place.
+    """
+    if loc.numel() == 0:
+        return
+    assert k_view.is_cuda and v_view.is_cuda, "store_cache_4d: CUDA only"
+    assert k_view.ndim == 4 and v_view.ndim == 4, (
+        f"store_cache_4d: k_view/v_view must be 4-D, "
+        f"got {k_view.ndim}/{v_view.ndim}"
+    )
+    assert cache_k.ndim == 3 and cache_v.ndim == 3, (
+        f"store_cache_4d: cache_k/cache_v must be 3-D, "
+        f"got {cache_k.ndim}/{cache_v.ndim}"
+    )
+    assert cache_k.shape[0] == cache_v.shape[0] == loc.numel(), (
+        "store_cache_4d: cache_k/cache_v/loc batch dim mismatch: "
+        f"{cache_k.shape[0]}, {cache_v.shape[0]}, {loc.numel()}"
+    )
+    assert k_view.dtype == v_view.dtype == cache_k.dtype == cache_v.dtype, (
+        "store_cache_4d: dtype mismatch: "
+        f"k_view={k_view.dtype}, v_view={v_view.dtype}, "
+        f"cache_k={cache_k.dtype}, cache_v={cache_v.dtype}"
+    )
+    # Stride invariants — the kernel addresses (head_num, head_dim) as one
+    # flat ROW_DIM dimension; this requires the trailing two dims to be
+    # contiguous. This holds for the page-major envelope views
+    # (k_stride = (page_bytes/itemsize, k_row_bytes/itemsize, head_dim, 1)) and
+    # for cache_k/cache_v produced by the model forward.
+    assert k_view.stride(-1) == 1 and k_view.stride(-2) == k_view.shape[-1], (
+        f"store_cache_4d: k_view trailing dims must be contiguous; "
+        f"got stride={k_view.stride()}, shape={tuple(k_view.shape)}"
+    )
+    assert v_view.stride(-1) == 1 and v_view.stride(-2) == v_view.shape[-1], (
+        f"store_cache_4d: v_view trailing dims must be contiguous; "
+        f"got stride={v_view.stride()}, shape={tuple(v_view.shape)}"
+    )
+    assert cache_k.stride(-1) == 1 and cache_k.stride(-2) == cache_k.shape[-1], (
+        f"store_cache_4d: cache_k trailing dims must be contiguous; "
+        f"got stride={cache_k.stride()}, shape={tuple(cache_k.shape)}"
+    )
+    assert cache_v.stride(-1) == 1 and cache_v.stride(-2) == cache_v.shape[-1], (
+        f"store_cache_4d: cache_v trailing dims must be contiguous; "
+        f"got stride={cache_v.stride()}, shape={tuple(cache_v.shape)}"
+    )
+
+    head_num = k_view.shape[2]
+    head_dim = k_view.shape[3]
+    v_head_dim = v_view.shape[3]
+    K_ROW_DIM = head_num * head_dim
+    V_ROW_DIM = head_num * v_head_dim
+    BLOCK = 128
+    N = loc.numel()
+    row_dim_max = max(K_ROW_DIM, V_ROW_DIM)
+    grid = (N, triton.cdiv(row_dim_max, BLOCK), 2)
+
+    store_cache_4d_kernel[grid](
+        k_view,
+        v_view,
+        cache_k,
+        cache_v,
+        loc,
+        k_view.stride(0),
+        k_view.stride(1),
+        v_view.stride(0),
+        v_view.stride(1),
+        cache_k.stride(0),
+        cache_v.stride(0),
+        K_ROW_DIM=K_ROW_DIM,
+        V_ROW_DIM=V_ROW_DIM,
+        PAGE_SIZE=page_size,
+        BLOCK=BLOCK,
+        num_warps=4,
+    )

python/sglang/kernels/ops/kvcache/cache_ops.py

@@ -0,0 +1,266 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def concat_and_cast_mha_k_kernel(
+    k_ptr,
+    k_nope_ptr,
+    k_rope_ptr,
+    head_cnt: tl.constexpr,
+    k_stride0: tl.constexpr,
+    k_stride1: tl.constexpr,
+    nope_stride0: tl.constexpr,
+    nope_stride1: tl.constexpr,
+    rope_stride0: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+):
+    pid_loc = tl.program_id(0)
+    head_range = tl.arange(0, head_cnt)
+
+    k_head_ptr = k_ptr + pid_loc * k_stride0 + head_range[:, None] * k_stride1
+
+    nope_offs = tl.arange(0, nope_dim)
+
+    src_nope_ptr = (
+        k_nope_ptr
+        + pid_loc * nope_stride0
+        + head_range[:, None] * nope_stride1
+        + nope_offs[None, :]
+    )
+    dst_nope_ptr = k_head_ptr + nope_offs[None, :]
+
+    src_nope = tl.load(src_nope_ptr)
+    tl.store(dst_nope_ptr, src_nope)
+
+    rope_offs = tl.arange(0, rope_dim)
+    src_rope_ptr = k_rope_ptr + pid_loc * rope_stride0 + rope_offs[None, :]
+    dst_rope_ptr = k_head_ptr + nope_dim + rope_offs[None, :]
+    src_rope = tl.load(src_rope_ptr)
+    tl.store(dst_rope_ptr, src_rope)
+
+
+def concat_and_cast_mha_k_triton(
+    k: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+):
+    # The source data type will be implicitly converted to the target data type.
+    assert (
+        len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3
+    ), f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+    assert (
+        k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0]
+    ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+    assert (
+        k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1]
+    ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+    assert (
+        k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1]
+    ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+
+    nope_dim = k_nope.shape[-1]
+    rope_dim = k_rope.shape[-1]
+    grid = (k.shape[0],)
+
+    concat_and_cast_mha_k_kernel[grid](
+        k,
+        k_nope,
+        k_rope,
+        k.shape[1],
+        k.stride(0),
+        k.stride(1),
+        k_nope.stride(0),
+        k_nope.stride(1),
+        k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+    )
+
+
+@triton.jit
+def reshape_and_cache_flash(
+    key_ptr,
+    value_ptr,
+    key_cache_ptr,
+    value_cache_ptr,
+    slot_mapping_ptr,
+    swa_slot_mapping_ptr,
+    k_scale_ptr,
+    v_scale_ptr,
+    block_stride,
+    key_stride,
+    value_stride,
+    num_heads,
+    head_size,
+    block_size,
+    HEAD_BLOCK: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+    USE_SCALE: tl.constexpr,
+):
+    """
+    Triton kernel for reshaping per-token K/V tensors into paged KV cache layout.
+
+    Source layout:
+        key/value: [num_tokens, num_heads, head_size]
+
+    Target cache layout:
+        cache: [num_blocks, block_size, num_heads, head_size]
+
+    Each Triton program instance handles:
+        - one token (program_id(0))
+        - one block of heads (program_id(1))
+
+    Features:
+        - optional SWA slot remapping
+        - optional FP8 scale dequantization before cache write
+
+    Args:
+        key_ptr: Pointer to source key tensor.
+        value_ptr: Pointer to source value tensor.
+        key_cache_ptr: Pointer to destination key cache tensor.
+        value_cache_ptr: Pointer to destination value cache tensor.
+        slot_mapping_ptr: Maps token -> cache slot.
+        swa_slot_mapping_ptr: Optional second-stage slot remap for SWA mode.
+        k_scale_ptr: Optional key scaling factor pointer.
+        v_scale_ptr: Optional value scaling factor pointer.
+        block_stride: Stride between cache blocks.
+        key_stride: Stride between source key tokens.
+        value_stride: Stride between source value tokens.
+        num_heads: Number of attention heads.
+        head_size: Hidden dimension per head.
+        block_size: Number of slots per cache block.
+        HEAD_BLOCK: Number of heads processed per program.
+        BLOCK_D: Vectorized dimension size (power-of-2 padded).
+        HAS_SWA: Enable SWA remapping.
+        USE_SCALE: Enable scale division before storing.
+    """
+
+    # ----------------------------------
+    # program ids
+    # pid0 = token
+    # pid1 = head block
+    # ----------------------------------
+    token_idx = tl.program_id(0)
+    head_block_idx = tl.program_id(1)
+
+    # ----------------------------------
+    # slot mapping
+    # ----------------------------------
+    slot_idx = tl.load(slot_mapping_ptr + token_idx)
+
+    if HAS_SWA:
+        slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx)
+
+    if slot_idx < 0:
+        return
+
+    block_idx = slot_idx // block_size
+    block_offset = slot_idx % block_size
+
+    # ----------------------------------
+    # head range
+    # ----------------------------------
+    head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK)
+
+    head_mask = head_idx < num_heads
+
+    dim_idx = tl.arange(0, BLOCK_D)
+
+    # shape = [HEAD_BLOCK, BLOCK_D]
+    offs = head_idx[:, None] * head_size + dim_idx[None, :]
+
+    mask = head_mask[:, None] & (dim_idx[None, :] < head_size)
+
+    # ----------------------------------
+    # source load
+    # ----------------------------------
+    src_key = token_idx * key_stride + offs
+    src_value = token_idx * value_stride + offs
+
+    k = tl.load(key_ptr + src_key, mask=mask)
+    v = tl.load(value_ptr + src_value, mask=mask)
+
+    # ----------------------------------
+    # optional scale
+    # ----------------------------------
+    if USE_SCALE:
+        k_scale = tl.load(k_scale_ptr)
+        v_scale = tl.load(v_scale_ptr)
+
+        k = k / k_scale
+        v = v / v_scale
+
+    # ----------------------------------
+    # target layout
+    # [block_idx, block_offset, head, dim]
+    # ----------------------------------
+    tgt = block_idx * block_stride + block_offset * num_heads * head_size + offs
+
+    tl.store(key_cache_ptr + tgt, k, mask=mask)
+    tl.store(value_cache_ptr + tgt, v, mask=mask)
+
+
+def launch_reshape_and_cache_flash(
+    key,
+    value,
+    key_cache,
+    value_cache,
+    slot_mapping,
+    swa_slot_mapping=None,
+    k_scale=None,
+    v_scale=None,
+):
+    """
+    Launch wrapper for reshape_and_cache_flash Triton kernel.
+
+    This wrapper prepares launch configuration and dispatches the Triton kernel
+    that writes token-major K/V tensors into paged KV cache layout.
+
+    Args:
+        key: Source key tensor [num_tokens, num_heads, head_size]
+        value: Source value tensor [num_tokens, num_heads, head_size]
+        key_cache: Destination key cache [num_blocks, block_size, num_heads, head_size]
+        value_cache: Destination value cache [num_blocks, block_size, num_heads, head_size]
+        slot_mapping: Token-to-cache slot mapping
+        swa_slot_mapping: Optional SWA remapping table
+        k_scale: Optional key scaling factor
+        v_scale: Optional value scaling factor
+    """
+
+    num_tokens = key.shape[0]
+    num_heads = key.shape[1]
+    head_size = key.shape[2]
+
+    HEAD_BLOCK = 4
+
+    BLOCK_D = triton.next_power_of_2(head_size)
+
+    grid = (
+        num_tokens,
+        triton.cdiv(num_heads, HEAD_BLOCK),
+    )
+
+    reshape_and_cache_flash[grid](
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        swa_slot_mapping,
+        k_scale if k_scale is not None else key,
+        v_scale if v_scale is not None else key,
+        key_cache.stride(0),
+        key.stride(0),
+        value.stride(0),
+        num_heads,
+        head_size,
+        key_cache.shape[1],
+        HEAD_BLOCK=HEAD_BLOCK,
+        BLOCK_D=BLOCK_D,
+        HAS_SWA=(swa_slot_mapping is not None),
+        USE_SCALE=(k_scale is not None),
+    )

python/sglang/kernels/ops/kvcache/kv_indices.py

@@ -0,0 +1,152 @@
+import triton
+import triton.language as tl
+
+_FLASHMLA_CREATE_KV_BLOCK_SIZE = 4096
+FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SIZE)
+
+
+@triton.jit
+def create_flashinfer_kv_indices_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices_ptr,
+    page_kernel_lens_ptr,
+    kv_indptr,
+    kv_start_idx,
+    kv_indices_ptr,
+    req_to_token_ptr_stride: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+    kv_indices_offset = tl.load(kv_indptr + pid)
+
+    kv_start = 0
+    kv_end = 0
+    if kv_start_idx:
+        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
+        kv_end = kv_start
+    kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
+
+    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+    for i in range(num_loop):
+        # index into req_to_token_ptr needs to be int64
+        offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE
+        mask = offset < kv_end - kv_start
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + kv_start
+            + offset,
+            mask=mask,
+        )
+        tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask)
+
+
+@triton.jit
+def create_chunked_prefix_cache_kv_indices(
+    req_to_token_ptr,  # (max_batch, max_context_len,)
+    req_pool_indices_ptr,  # (batch_size,)
+    chunk_start_idx_ptr,  # (batch_size,)
+    chunk_seq_lens_ptr,  # (batch_size,)
+    chunk_cu_seq_lens_ptr,  # (batch_size + 1,)
+    chunk_kv_indices_ptr,  # (num_chunk_tokens,)
+    req_to_token_ptr_stride: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+    chunk_kv_indices_offset = tl.load(chunk_cu_seq_lens_ptr + pid)
+
+    # get the token positions of current chunk
+    chunk_start_pos = tl.load(chunk_start_idx_ptr + pid).to(tl.int32)
+    chunk_seq_len = tl.load(chunk_seq_lens_ptr + pid).to(tl.int32)
+
+    num_loop = tl.cdiv(chunk_seq_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < chunk_seq_len
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + chunk_start_pos
+            + offset,
+            mask=mask,
+        )
+        tl.store(
+            chunk_kv_indices_ptr + chunk_kv_indices_offset + offset, data, mask=mask
+        )
+
+
+def get_num_page_per_block_flashmla(page_size: int = 64) -> int:
+    num_page_per_block = _FLASHMLA_CREATE_KV_BLOCK_SIZE // page_size
+    return num_page_per_block
+
+
+def get_num_kv_index_blocks_flashmla(kv_indices_width: int, page_size: int) -> int:
+    """Grid axis-1 size for create_flashmla_kv_indices_triton: the number of
+    page-blocks spanning the widest sequence (one CTA per block). kv_indices_width
+    is the per-row width of the kv_indices buffer (the kernel's kv_indices_ptr_stride).
+    """
+    npb = get_num_page_per_block_flashmla(page_size)
+    return (kv_indices_width + npb - 1) // npb
+
+
+@triton.jit
+def create_flashmla_kv_indices_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices_ptr,
+    page_kernel_lens_ptr,
+    kv_start_idx,
+    kv_indices_ptr,
+    req_to_token_ptr_stride: tl.constexpr,
+    kv_indices_ptr_stride: tl.constexpr,
+    PAGED_SIZE: tl.constexpr = 64,
+):
+    NUM_PAGE_PER_BLOCK: tl.constexpr = (
+        FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE
+    )
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+
+    kv_start = 0
+    kv_end = 0
+    if kv_start_idx:
+        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
+        kv_end = kv_start
+
+    kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
+
+    num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE)
+    num_pages_loop = tl.cdiv(kv_end - kv_start, FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON)
+
+    # One CTA per page-block (grid axis 1) rather than one CTA looping all blocks;
+    # CTAs beyond this sequence's block count are guarded out.
+    i = tl.program_id(axis=1)
+    if i < num_pages_loop:
+        # index into req_to_token_ptr needs to be int64
+        paged_offset = (
+            tl.arange(0, NUM_PAGE_PER_BLOCK).to(tl.int64) + i * NUM_PAGE_PER_BLOCK
+        ) * PAGED_SIZE
+        paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK
+
+        mask = paged_offset < num_paged * PAGED_SIZE
+        mask_out = paged_offset_out < num_paged
+
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + kv_start
+            + paged_offset,
+            mask=mask,
+        )
+        tl.store(
+            kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out,
+            data // PAGED_SIZE,
+            mask=mask_out,
+        )

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.utils import is_arch_support_pdl
+from sglang.srt.runtime_context import get_parallel
+
+
+@triton.jit
+def set_mla_kv_buffer_kernel(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    BLOCK: tl.constexpr,
+    DCP_RANK: tl.constexpr,
+    DCP_WORLD_SIZE: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
+):
+    pid_loc = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    base = pid_blk * BLOCK
+    offs = base + tl.arange(0, BLOCK)
+    total_dim = nope_dim + rope_dim
+    mask = offs < total_dim
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    is_valid = loc % DCP_WORLD_SIZE == DCP_RANK
+    safe_loc = tl.where(is_valid, loc, 0)
+    safe_loc = safe_loc // DCP_WORLD_SIZE
+    dst_ptr = kv_buffer_ptr + safe_loc * buffer_stride + offs
+
+    # Three-way branch to handle boundary correctly while preserving fast path
+    if base + BLOCK <= nope_dim:
+        # Fast path: entire block is in nope region
+        src = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask,
+        )
+    elif base >= nope_dim:
+        # Fast path: entire block is in rope region
+        offs_rope = offs - nope_dim
+        src = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + offs_rope,
+            mask=mask,
+        )
+    else:
+        # Boundary case: block spans nope/rope boundary (e.g., FP8 with nope_dim=528)
+        # Handle each offset individually to avoid negative indexing
+        is_nope = offs < nope_dim
+        is_rope = (offs >= nope_dim) & (offs < (nope_dim + rope_dim))
+
+        src_nope = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask & is_nope,
+            other=0,
+        )
+        src_rope = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim),
+            mask=mask & is_rope,
+            other=0,
+        )
+
+        src = tl.where(is_nope, src_nope, src_rope)
+
+    tl.store(dst_ptr, src, mask=mask & is_valid)
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+# Above this loc count the TMA bulk-store path overtakes the single-CTA-per-loc
+# Triton kernel. Below it, Triton with BLOCK = next_pow2(total_dim) (one CTA
+# does the whole row in one tile, no boundary fan-out) is the winning fallback.
+# Tuned on GB300 with DSv4 row widths.
+_TMA_BULK_STORE_MIN_LOCS = 768
+
+
+def set_mla_kv_buffer_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+):
+    """Dispatch MLA paged-KV scatter writes to the fastest available path.
+
+    Two paths, chosen on ``n_loc``:
+
+    - ``n_loc >= 768`` (and SM90+ with TMA-compatible row widths): JIT CUDA
+      kernel where each warp loads one (nope, rope) row into shared memory and
+      issues a single ``cp.async.bulk.global.shared::cta`` store to scatter the
+      row at ``kv_buffer[loc[item]]``. Wins at large bs because it packs 4-8
+      items per CTA, drastically reducing the CTA count vs single-CTA-per-loc.
+    - Otherwise: Triton kernel with ``BLOCK = next_pow2(nope_dim + rope_dim)``,
+      i.e. one CTA per loc covering the entire row in one tile. Wins at small
+      bs because there's no per-loc CTA fan-out (5x fewer CTAs than the old
+      BLOCK=128 dispatch) and the row-spanning block makes the boundary branch
+      a one-shot per CTA. This is also the path for SM<90 and for shapes that
+      violate the TMA 16-byte alignment.
+
+    Speedup vs the legacy BLOCK=128 Triton kernel on GB300 (BF16, nope=512,
+    rope=64): ~1.05x at bs=8, ~1.5x at bs=128, 3.5x at bs=512, **11.7x at
+    bs=16384**.
+
+    Name retained for caller compatibility; the implementation is no longer
+    Triton-only.
+    """
+    from sglang.jit_kernel.set_mla_kv_buffer import (
+        can_use_set_mla_kv_buffer,
+    )
+    from sglang.jit_kernel.set_mla_kv_buffer import (
+        set_mla_kv_buffer as jit_set_mla_kv_buffer,
+    )
+
+    n_loc = loc.numel()
+    nope_bytes = cache_k_nope.shape[-1] * cache_k_nope.element_size()
+    rope_bytes = cache_k_rope.shape[-1] * cache_k_rope.element_size()
+    if (
+        n_loc >= _TMA_BULK_STORE_MIN_LOCS
+        and is_arch_support_pdl()
+        and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes)
+        and not get_parallel().dcp_enabled
+    ):
+        jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
+        return
+
+    # Fallback: Triton with BLOCK = next_pow2(total_dim). One CTA per loc; the
+    # whole row in one tile (the existing 3-way nope/rope/boundary branch in
+    # ``set_mla_kv_buffer_kernel`` handles the over-allocation past total_dim
+    # via the offs<total_dim mask). Beats BLOCK=128 by 60-2700 ns across the
+    # 2 <= bs <= 512 range on GB300.
+    nope_dim = cache_k_nope.shape[-1]
+    rope_dim = cache_k_rope.shape[-1]
+    total_dim = nope_dim + rope_dim
+    BLOCK = triton.next_power_of_2(total_dim)
+    grid = (n_loc, 1)
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+    set_mla_kv_buffer_kernel[grid](
+        kv_buffer,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+        BLOCK=BLOCK,
+        DCP_RANK=get_parallel().attn_dcp_rank,
+        DCP_WORLD_SIZE=get_parallel().attn_dcp_size,
+        **pdl_kwargs,
+    )
+
+
+@triton.jit
+def set_mla_kv_buffer_fp8_quant_kernel(
+    kv_buffer_fp8_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    BLOCK: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
+):
+    """Fuse BF16/FP16->FP8 cast with paged KV write."""
+    pid_loc = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    base = pid_blk * BLOCK
+    offs = base + tl.arange(0, BLOCK)
+    total_dim = nope_dim + rope_dim
+    mask = offs < total_dim
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs
+
+    if base + BLOCK <= nope_dim:
+        src = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask,
+            other=0.0,
+        )
+    elif base >= nope_dim:
+        offs_rope = offs - nope_dim
+        src = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + offs_rope,
+            mask=mask,
+            other=0.0,
+        )
+    else:
+        is_nope = offs < nope_dim
+        src_nope = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask & is_nope,
+            other=0.0,
+        )
+        src_rope = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim),
+            mask=mask & ~is_nope,
+            other=0.0,
+        )
+        src = tl.where(is_nope, src_nope, src_rope)
+
+    # Destination pointer is FP8-typed view; tl.store performs downcast.
+    tl.store(dst_ptr, src, mask=mask)
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def set_mla_kv_buffer_triton_fp8_quant(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+    fp8_dtype: torch.dtype,
+):
+    """Fuse BF16/FP16 MLA K quantization with paged KV write."""
+    kv_buffer_fp8 = kv_buffer.view(fp8_dtype)
+
+    nope_dim = cache_k_nope.shape[-1]
+    rope_dim = cache_k_rope.shape[-1]
+    total_dim = nope_dim + rope_dim
+    BLOCK = 128
+    n_loc = loc.numel()
+    grid = (n_loc, triton.cdiv(total_dim, BLOCK))
+
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+
+    set_mla_kv_buffer_fp8_quant_kernel[grid](
+        kv_buffer_fp8,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer_fp8.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+        BLOCK=BLOCK,
+        **pdl_kwargs,
+    )
+
+
+@triton.jit
+def set_mla_kv_scale_buffer_kernel(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    pid_loc = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    base = pid_blk * BLOCK
+    offs = base + tl.arange(0, BLOCK)
+    total_dim = nope_dim + rope_dim
+    mask = offs < total_dim  # Make sure don't cross the boundary
+
+    loc = tl.load(loc_ptr + pid_loc)
+    dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
+
+    # Check each offs should read 'nope' or 'rope'
+    is_nope = offs < nope_dim
+    src_nope = tl.load(
+        cache_k_nope_ptr + pid_loc * nope_stride + offs, mask=mask & is_nope, other=0.0
+    )
+    src_rope = tl.load(
+        cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim),
+        mask=mask & ~is_nope,
+        other=0.0,
+    )
+
+    # Combine nope + rope
+    src = src_nope + src_rope
+    tl.store(dst_ptr, src, mask=mask)
+
+
+def set_mla_kv_scale_buffer_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+):
+    nope_dim = cache_k_nope.shape[-1]
+    rope_dim = cache_k_rope.shape[-1]
+    total_dim = nope_dim + rope_dim
+    BLOCK = 128  # Keep origin, works for smaller total_dim as well.
+    n_loc = loc.numel()
+    grid = (n_loc, triton.cdiv(total_dim, BLOCK))
+
+    set_mla_kv_scale_buffer_kernel[grid](
+        kv_buffer,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+        BLOCK=BLOCK,
+    )
+
+
+@triton.jit
+def get_mla_kv_buffer_kernel(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+):
+    pid_loc = tl.program_id(0)
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    loc_src_ptr = kv_buffer_ptr + loc * buffer_stride
+
+    nope_offs = tl.arange(0, nope_dim)
+    nope_src_ptr = loc_src_ptr + nope_offs
+    nope_src = tl.load(nope_src_ptr)
+
+    tl.store(
+        cache_k_nope_ptr + pid_loc * nope_stride + nope_offs,
+        nope_src,
+    )
+
+    rope_offs = tl.arange(0, rope_dim)
+    rope_src_ptr = loc_src_ptr + nope_dim + rope_offs
+    rope_src = tl.load(rope_src_ptr)
+    tl.store(
+        cache_k_rope_ptr + pid_loc * rope_stride + rope_offs,
+        rope_src,
+    )
+
+
+def get_mla_kv_buffer_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+):
+    # The source data type will be implicitly converted to the target data type.
+    nope_dim = cache_k_nope.shape[-1]  # 512
+    rope_dim = cache_k_rope.shape[-1]  # 64
+    n_loc = loc.numel()
+    grid = (n_loc,)
+
+    get_mla_kv_buffer_kernel[grid](
+        kv_buffer,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+    )

python/sglang/kernels/ops/kvcache/rope_cache.py

@@ -0,0 +1,736 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _get_gptj_rotated_x(
+    x,
+    x_rotated_mask,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    # GPT-J rotary layout:
+    # Pair adjacent dimensions and apply:
+    # [x0, x1, x2, x3] -> [-x1, x0, -x3, x2]
+
+    # Apply sign inversion on odd positions.
+    x_rotated = tl.where(x_rotated_mask, x, -x)
+    # Reshape into (D/2, 2) pairs.
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D_HALF, 2))
+    # Swap each pair.
+    x_rotated = tl.flip(x_rotated, 1)
+    # Flatten back to original shape.
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
+    return x_rotated
+
+
+@triton.jit
+def _get_neox_rotated_x(
+    x,
+    x_rotated_mask,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    # GPT-NeoX rotary layout:
+    # Split head dimension into two halves:
+    # [x0, x1, x2, x3] -> [-x2, -x3, x0, x1]
+
+    # Keep first half positive, second half negative.
+    x_rotated = tl.where(x_rotated_mask, x, -x)
+    # Reshape into (2, D/2).
+    x_rotated = tl.reshape(x_rotated, (2, BLOCK_D_HALF))
+    # Reverse each half.
+    x_rotated = tl.flip(x_rotated, 1)
+    # Flatten and reverse full vector.
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
+    x_rotated = tl.flip(x_rotated, 0)
+    return x_rotated
+
+
+@triton.jit
+def _unit_rope(
+    x_ptrs,
+    cos,
+    sin,
+    d_pe_offs,
+    IS_NEOX: tl.constexpr,
+    BLOCK_D_pe: tl.constexpr,
+    BLOCK_D_HALF_pe: tl.constexpr,
+):
+    # Load one full attention head vector.
+    x_pe = tl.load(x_ptrs)
+
+    # Stage 1: Build rotated vector according to rotary layout.
+    if IS_NEOX:
+        x_rotated_mask = d_pe_offs < BLOCK_D_HALF_pe
+        x_pe_rotated = _get_neox_rotated_x(
+            x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe
+        )
+    else:
+        x_rotated_mask = d_pe_offs % 2 == 0
+        x_pe_rotated = _get_gptj_rotated_x(
+            x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe
+        )
+
+    # Stage 2: Apply RoPE transform:
+    # x' = x*cos + rotate(x)*sin
+    x_pe = x_pe * cos + x_pe_rotated * sin
+
+    return x_pe
+
+
+@triton.jit
+def _load_cos_sin(
+    cos_sin_ptr,
+    pos,
+    d_cos_offs,
+    stride_t,
+    stride_d,
+    freq_dim,
+):
+    base = pos * stride_t
+    cos = tl.load(cos_sin_ptr + base + d_cos_offs * stride_d)
+    sin = tl.load(cos_sin_ptr + base + (d_cos_offs + freq_dim) * stride_d)
+    return cos, sin
+
+
+@triton.jit
+def _fused_qk_rope_reshape_and_cache_kernel(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    pos_ptr,
+    cos_sin_ptr,
+    offs_ptr,
+    key_cache_ptr,
+    value_cache_ptr,
+    slot_mapping_ptr,
+    swa_slot_mapping_ptr,
+    q_out_ptr,
+    k_out_ptr,
+    zeros_out_ptr,
+    T,
+    T_slot,
+    q_stride_t,
+    q_stride_h,
+    q_stride_d,
+    k_stride_t,
+    k_stride_h,
+    k_stride_d,
+    v_stride_t,
+    v_stride_h,
+    v_stride_d,
+    cos_sin_stride_t,
+    cos_sin_stride_d,
+    q_out_stride_t,
+    q_out_stride_h,
+    q_out_stride_d,
+    k_out_stride_t,
+    k_out_stride_h,
+    k_out_stride_d,
+    key_cache_stride_t,
+    key_cache_stride_h,
+    key_cache_stride_d,
+    key_cache_stride_b,
+    key_cache_stride_x,
+    value_cache_stride_t,
+    value_cache_stride_h,
+    value_cache_stride_d,
+    value_cache_stride_b,
+    value_cache_stride_slot_chunk,
+    value_cache_stride_x,
+    zeros_out_stride_t,
+    zeros_out_stride_h,
+    zeros_out_stride_d,
+    k_scale_ptr,
+    v_scale_ptr,
+    QH_PER_KH: tl.constexpr,
+    QH: tl.constexpr,
+    KH: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    BLOCK_D_pe: tl.constexpr,
+    BLOCK_D_HALF_pe: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    X_SIZE: tl.constexpr,
+    FLASH_LAYOUT: tl.constexpr,
+    VALUE_SHUFFLE_LAYOUT: tl.constexpr = False,
+    HAVE_POS: tl.constexpr = False,
+    HAVE_K_SCALE: tl.constexpr = False,
+    HAVE_V_SCALE: tl.constexpr = False,
+    HAVE_ZEROS: tl.constexpr = False,
+    HAS_SWA: tl.constexpr = False,
+):
+    # ============================================================
+    # Stage 0: Static stride assumptions for Triton compiler
+    #
+    # These assumptions help Triton optimize pointer arithmetic and
+    # simplify generated address calculations.
+    # ============================================================
+
+    tl.assume(q_stride_t >= 0)
+    tl.assume(q_stride_h >= 0)
+    tl.assume(q_stride_d >= 0)
+    tl.assume(k_stride_t >= 0)
+    tl.assume(k_stride_h >= 0)
+    tl.assume(k_stride_d >= 0)
+    tl.assume(v_stride_t >= 0)
+    tl.assume(v_stride_h >= 0)
+    tl.assume(v_stride_d >= 0)
+    tl.assume(cos_sin_stride_t >= 0)
+    tl.assume(cos_sin_stride_d >= 0)
+    tl.assume(q_out_stride_t >= 0)
+    tl.assume(q_out_stride_h >= 0)
+    tl.assume(q_out_stride_d >= 0)
+    tl.assume(k_out_stride_t >= 0)
+    tl.assume(k_out_stride_h >= 0)
+    tl.assume(k_out_stride_d >= 0)
+    tl.assume(key_cache_stride_t >= 0)
+    tl.assume(key_cache_stride_h >= 0)
+    tl.assume(key_cache_stride_d >= 0)
+    tl.assume(key_cache_stride_b >= 0)
+    tl.assume(key_cache_stride_x >= 0)
+    tl.assume(value_cache_stride_t >= 0)
+    tl.assume(value_cache_stride_h >= 0)
+    tl.assume(value_cache_stride_d >= 0)
+    tl.assume(value_cache_stride_b >= 0)
+    tl.assume(value_cache_stride_slot_chunk >= 0)
+    tl.assume(value_cache_stride_x >= 0)
+    tl.assume(zeros_out_stride_t >= 0)
+    tl.assume(zeros_out_stride_h >= 0)
+    tl.assume(zeros_out_stride_d >= 0)
+
+    # ============================================================
+    # Stage 1: Program instance mapping
+    #
+    # Each program handles:
+    #   - one (token, q_head) for Q path
+    #   - selected KV ownership for cache write path
+    #
+    # pid layout:
+    #   [0, T*QH)            -> decode Q path
+    #   [T*QH, extra KV)     -> KV-only path
+    # ============================================================
+
+    pid = tl.program_id(0)
+    tl.assume(pid >= 0)
+
+    d_pe_offs = tl.arange(0, BLOCK_D_pe).to(tl.int64)
+
+    # ============================================================
+    # Stage 2: Main decode path (Q always active)
+    # ============================================================
+
+    if pid < T * QH:
+        pid_t = pid // QH
+        pid_hq = pid % QH
+
+        # --------------------------------------------------------
+        # Stage 2.1: Compute rotary frequency offsets
+        #
+        # RoPE frequencies may be stored as:
+        #   D/2 frequencies (shared front-half)
+        #   D frequencies (full explicit)
+        # --------------------------------------------------------
+
+        if REUSE_FREQS_FRONT_PART:
+            if IS_NEOX:
+                d_cos_offs = d_pe_offs
+                d_cos_offs = tl.where(
+                    (d_cos_offs >= BLOCK_D_HALF_pe) & (d_cos_offs < BLOCK_D_pe),
+                    d_cos_offs - BLOCK_D_HALF_pe,
+                    d_cos_offs,
+                ).to(d_cos_offs.dtype)
+                # d_cos_mask = d_cos_offs < BLOCK_D_pe
+            else:
+                d_cos_offs = d_pe_offs // 2
+                # d_cos_mask = d_cos_offs < BLOCK_D_HALF_pe
+        else:
+            d_cos_offs = d_pe_offs
+            # d_cos_mask = d_cos_offs < BLOCK_D_pe
+
+        # --------------------------------------------------------
+        # Stage 2.2: Load token position and optional offset
+        #
+        # offs_ptr is used by chunked prefill / sliding-window decode.
+        # --------------------------------------------------------
+        pos = tl.load(pos_ptr + pid_t)
+        if HAVE_POS:
+            offset = tl.load(offs_ptr + pid_t)
+            pos = pos + offset
+
+        # --------------------------------------------------------
+        # Stage 2.3: Load cosine / sine table
+        # --------------------------------------------------------
+        # cos_offs = pos * cos_stride_t + d_cos_offs * cos_stride_d
+        # cos = tl.load(cos_ptr + cos_offs)
+        # sin = tl.load(sin_ptr + cos_offs)
+
+        freq_dim = BLOCK_D_HALF_pe if REUSE_FREQS_FRONT_PART else BLOCK_D_pe
+
+        cos, sin = _load_cos_sin(
+            cos_sin_ptr,
+            pos,
+            d_cos_offs,
+            cos_sin_stride_t,
+            cos_sin_stride_d,
+            freq_dim,
+        )
+
+        # --------------------------------------------------------
+        # Stage 2.4: Apply RoPE to Q
+        # --------------------------------------------------------
+        q_ptrs = (
+            q_ptr + pid_t * q_stride_t + pid_hq * q_stride_h + d_pe_offs * q_stride_d
+        )
+        q_pe = _unit_rope(
+            q_ptrs,
+            cos,
+            sin,
+            d_pe_offs,
+            IS_NEOX,
+            BLOCK_D_pe,
+            BLOCK_D_HALF_pe,
+        )
+
+        # Store rotated Q output.
+        q_out_ptrs = (
+            q_out_ptr
+            + pid_t * q_out_stride_t
+            + pid_hq * q_out_stride_h
+            + d_pe_offs * q_out_stride_d
+        )
+        tl.store(q_out_ptrs, q_pe.to(q_out_ptr.dtype.element_ty))
+
+        if HAVE_ZEROS:
+            z = tl.zeros((BLOCK_D_pe,), dtype=zeros_out_ptr.dtype.element_ty)
+            zeros_out_ptrs = (
+                zeros_out_ptr
+                + pid_t * zeros_out_stride_t
+                + pid_hq * zeros_out_stride_h
+                + d_pe_offs * zeros_out_stride_d
+            )
+            tl.store(zeros_out_ptrs, z)
+
+        # ========================================================
+        # Stage 3: KV ownership path
+        #
+        # Only one Q group leader writes KV:
+        #   pid_hq % QH_PER_KH == 0
+        #
+        # This prevents duplicated KV cache writes.
+        # ========================================================
+
+        if pid_hq % QH_PER_KH == 0:
+            # ----------------------------------------------------
+            # Stage 3.1: Resolve cache slot
+            # ----------------------------------------------------
+            pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64)
+            if HAS_SWA:
+                pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot)
+
+            # ------------------------------------------------
+            # Stage 3.2: Apply RoPE to K
+            # ------------------------------------------------
+            if pid_slot >= 0:
+                pid_t_slot = pid_slot // BLOCK_SIZE
+                pid_b = pid_slot % BLOCK_SIZE
+                pid_hk = pid_hq // QH_PER_KH
+                if HAVE_K_SCALE:
+                    k_scale = tl.load(k_scale_ptr)
+                else:
+                    k_scale = 1
+                k_ptrs = (
+                    k_ptr
+                    + pid_t * k_stride_t
+                    + pid_hk * k_stride_h
+                    + d_pe_offs * k_stride_d
+                )
+                k_pe = _unit_rope(
+                    k_ptrs,
+                    cos,
+                    sin,
+                    d_pe_offs,
+                    IS_NEOX,
+                    BLOCK_D_pe,
+                    BLOCK_D_HALF_pe,
+                )
+
+                k_out_ptrs = (
+                    k_out_ptr
+                    + pid_t * k_out_stride_t
+                    + pid_hk * k_out_stride_h
+                    + d_pe_offs * k_out_stride_d
+                )
+                tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty))
+
+                # ------------------------------------------------
+                # Stage 3.3: Optional fp8 scaling before cache
+                # ------------------------------------------------
+
+                k_scale_rcprl = 1 / k_scale
+                k_pe = k_pe * k_scale_rcprl
+
+                # ------------------------------------------------
+                # Stage 3.4: Write K cache
+                #
+                # Two layouts supported:
+                #   FLASH_LAYOUT
+                #   paged KV layout
+                # ------------------------------------------------
+
+                if FLASH_LAYOUT:
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + pid_b * key_cache_stride_b
+                        + pid_hk * key_cache_stride_h
+                        + d_pe_offs * key_cache_stride_d
+                    )
+                else:
+                    k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE))
+                    dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64)
+                    x_offs = tl.arange(0, X_SIZE).to(tl.int64)
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + pid_hk * key_cache_stride_h
+                        + dx_offs[:, None] * key_cache_stride_d
+                        + pid_b * key_cache_stride_b
+                        + x_offs[None, :] * key_cache_stride_x
+                    )
+
+                tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty))
+
+                # ------------------------------------------------
+                # Stage 3.5: Write V cache
+                #
+                # Supports:
+                #   normal layout
+                #   shuffle layout
+                # ------------------------------------------------
+
+                v_ptrs = (
+                    v_ptr
+                    + pid_t * v_stride_t
+                    + pid_hk * v_stride_h
+                    + d_pe_offs * v_stride_d
+                )
+                if HAVE_V_SCALE:
+                    v_scale = tl.load(v_scale_ptr)
+                else:
+                    v_scale = 1
+                v_scale_rcprl = 1 / v_scale
+                v = tl.load(v_ptrs) * v_scale_rcprl
+                if VALUE_SHUFFLE_LAYOUT:
+                    slot_chunk = pid_b // X_SIZE
+                    x_off = pid_b % X_SIZE
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + slot_chunk * value_cache_stride_slot_chunk
+                        + d_pe_offs.to(tl.int64) * value_cache_stride_d
+                        + x_off * value_cache_stride_x
+                    )
+                else:
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + d_pe_offs.to(tl.int64) * value_cache_stride_d
+                        + pid_b * value_cache_stride_b
+                    )
+                tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty))
+    # ============================================================
+    # Stage 4: Extra KV-only path
+    #
+    # Handles tokens that only require cache update:
+    #   T_slot > T
+    #
+    # No Q / no RoPE on Q branch.
+    # ============================================================
+    else:
+        pid = pid - T * QH + T * KH
+        if pid < T_slot * KH:
+            pid_t = pid // KH
+            pid_hk = pid % KH
+            pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64)
+            if HAS_SWA:
+                pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot)
+
+            if pid_slot >= 0:
+                pid_t_slot = pid_slot // BLOCK_SIZE
+                pid_b = pid_slot % BLOCK_SIZE
+                if HAVE_K_SCALE:
+                    k_scale = tl.load(k_scale_ptr)
+                else:
+                    k_scale = 1
+                k_ptrs = (
+                    k_ptr
+                    + pid_t * k_stride_t
+                    + pid_hk * k_stride_h
+                    + d_pe_offs * k_stride_d
+                )
+
+                k_pe = tl.load(k_ptrs)
+
+                k_out_ptrs = (
+                    k_out_ptr
+                    + pid_t * k_out_stride_t
+                    + pid_hk * k_out_stride_h
+                    + d_pe_offs * k_out_stride_d
+                )
+                tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty))
+
+                k_scale_rcprl = 1 / k_scale
+                k_pe = k_pe * k_scale_rcprl
+
+                if FLASH_LAYOUT:
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + d_pe_offs * key_cache_stride_d
+                        + pid_b * key_cache_stride_b
+                        + pid_hk * key_cache_stride_h
+                    )
+                else:
+                    k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE))
+                    dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64)
+                    x_offs = tl.arange(0, X_SIZE).to(tl.int64)
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + pid_hk * key_cache_stride_h
+                        + dx_offs[:, None] * key_cache_stride_d
+                        + pid_b * key_cache_stride_b
+                        + x_offs[None, :] * key_cache_stride_x
+                    )
+                tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty))
+
+                v_ptrs = (
+                    v_ptr
+                    + pid_t * v_stride_t
+                    + pid_hk * v_stride_h
+                    + d_pe_offs * v_stride_d
+                )
+                if HAVE_V_SCALE:
+                    v_scale = tl.load(v_scale_ptr)
+                else:
+                    v_scale = 1
+                v_scale_rcprl = 1 / v_scale
+                v = tl.load(v_ptrs) * v_scale_rcprl
+                if VALUE_SHUFFLE_LAYOUT:
+                    slot_chunk = pid_b // X_SIZE
+                    x_off = pid_b % X_SIZE
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + slot_chunk * value_cache_stride_slot_chunk
+                        + d_pe_offs * value_cache_stride_d
+                        + x_off * value_cache_stride_x
+                    )
+                else:
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + d_pe_offs * value_cache_stride_d
+                        + pid_b * value_cache_stride_b
+                    )
+                tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty))
+
+
+def fused_qk_rope_reshape_and_cache(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    pos: torch.Tensor,
+    cos_sin: torch.Tensor,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+    is_neox: bool,
+    flash_layout: bool,
+    apply_scale: bool = True,
+    offs: torch.Tensor = None,
+    q_out: torch.Tensor = None,
+    k_out: torch.Tensor = None,
+    output_zeros: bool = True,
+    zeros_out: torch.Tensor = None,
+    swa_slot_mapping=None,
+):
+    """
+    Perform RoPE on q and k and along the last dimension and copy k and v in to key_cache and value_cache inplace
+
+    Key parameters:
+    - q: shape (T, QH, D).
+    - k: shape (T_slot, KH, D).
+    - v: shape (T_slot, KH, D).
+    - if flash_layout:
+    -     key_cache: shape (T_cache, block_size, KH, D).
+    -     value_cache: shape (T_cache, block_size, KH, D).
+    - else:
+    -     key_cache: shape (T_cache, KH, D // x, block_size, x).
+    -     value_cache: shape (T_cache, KH, D, block_size).
+    - slot_mapping: shape (T_slot, ).
+
+    T is the number of decode tokens, T_cahce * block_size is the max number of tokens of kv_cache
+    QH must be multiple of KH
+
+    Returns:
+    - q_out: same shape as input q.
+    - k_out: same shape as input k.
+    - key_cache: same shape as input key_cache (inplace).
+    - value_cache: same shape as input value_cache (inplace).
+    - zeros_out: same shape as input q.
+    """
+
+    t, qh, d = q.shape
+    tk, kh, dk = k.shape
+    tv, vh, dv = v.shape
+    if flash_layout:
+        t_cache, block_size, kh_cache, dk_cache = key_cache.shape
+        t_cache_v, block_size_v, vh_cache, dv_cache = value_cache.shape
+        value_shuffle_layout = False
+    else:
+        t_cache, kh_cache, dkx_cache, block_size, x_cache = key_cache.shape
+        if value_cache.ndim == 5:
+            # value_cache shuffle: (num_blocks, num_kv_heads, block_size // x, head_size, x)
+            t_cache_v, vh_cache, slot_chunk_v, dv_cache, x_v = value_cache.shape
+            value_shuffle_layout = True
+            block_size_v = slot_chunk_v * x_v
+            assert block_size_v == block_size and x_v == x_cache, (
+                f"value_cache shuffle (T,KH,block_size//x,D,x) must match key: "
+                f"{block_size_v=} {block_size=} {x_v=} {x_cache=}"
+            )
+        else:
+            t_cache_v, vh_cache, dv_cache, block_size_v = value_cache.shape
+            value_shuffle_layout = False
+    (t_slot,) = slot_mapping.shape
+
+    assert (
+        t == tk == tv and t_slot <= tk
+    ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}"
+    assert (
+        block_size == block_size_v
+    ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}"
+    assert (
+        kh == vh == kh_cache == vh_cache
+    ), "KV head should be identical for k, v, key_cache, and value_cache"
+    assert (
+        t_cache == t_cache_v
+    ), "Number of tokens should be identical for key_cache, and value_cache"
+    if flash_layout:
+        assert (
+            d == dk == dv == dk_cache == dv_cache
+        ), "D dimension should be identical for q, k, and v"
+    else:
+        assert (
+            d == dk == dv == dkx_cache * x_cache == dv_cache
+        ), "D dimension should be identical for q, k, and v"
+        assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2"
+
+    assert d == triton.next_power_of_2(d), "D dimension should be power of 2"
+    assert block_size == triton.next_power_of_2(
+        block_size
+    ), "block_size should be power of 2"
+    assert qh % kh == 0, "Q heads must be multiple of H heads"
+    d_freq = cos_sin.shape[-1] // 2
+    assert (d_freq == d // 2) or (
+        d_freq == d
+    ), "cos/sin last dim should be the same or half of the qk last dim"
+    reuse_freqs_front_part = d_freq == d // 2
+
+    if q_out is None:
+        q_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)
+
+    if k_out is None:
+        k_out = torch.empty((tk, kh, dk), dtype=k.dtype, device=q.device)
+
+    if zeros_out is not None:
+        tz, qhz, dz = zeros_out.shape
+        assert (
+            t == tz and qh == qhz and d == dz
+        ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}"
+        output_zeros = True
+    elif output_zeros:
+        zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)
+    else:
+        zeros_out = None
+
+    n_pid = t * qh + (t_slot - t) * kh if t_slot >= t else t * qh
+    grid = (n_pid, 1, 1)
+    _fused_qk_rope_reshape_and_cache_kernel[grid](
+        q,
+        k,
+        v,
+        pos,
+        cos_sin,
+        offs,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        swa_slot_mapping,
+        q_out,
+        k_out,
+        zeros_out,
+        t,
+        t_slot,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        cos_sin.stride(0),
+        cos_sin.stride(-1),
+        *q_out.stride(),
+        *k_out.stride(),
+        key_cache.stride(0) if not flash_layout else key_cache.stride(0),
+        key_cache.stride(1) if not flash_layout else key_cache.stride(2),
+        key_cache.stride(2) if not flash_layout else key_cache.stride(3),
+        key_cache.stride(3) if not flash_layout else key_cache.stride(1),
+        key_cache.stride(4) if not flash_layout else 0,
+        value_cache.stride(0) if not flash_layout else value_cache.stride(0),
+        value_cache.stride(1) if not flash_layout else value_cache.stride(2),
+        (
+            value_cache.stride(3)
+            if (not flash_layout and value_shuffle_layout)
+            else (value_cache.stride(2) if not flash_layout else value_cache.stride(3))
+        ),
+        (
+            0
+            if (not flash_layout and value_shuffle_layout)
+            else (value_cache.stride(3) if not flash_layout else value_cache.stride(1))
+        ),
+        value_cache.stride(2) if (not flash_layout and value_shuffle_layout) else 0,
+        value_cache.stride(4) if (not flash_layout and value_shuffle_layout) else 0,
+        zeros_out.stride(0) if zeros_out is not None else 0,
+        zeros_out.stride(1) if zeros_out is not None else 0,
+        zeros_out.stride(2) if zeros_out is not None else 0,
+        k_scale_ptr=k_scale,
+        v_scale_ptr=v_scale,
+        QH_PER_KH=qh // kh,
+        QH=qh,
+        KH=kh,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=is_neox,
+        BLOCK_D_pe=d,
+        BLOCK_D_HALF_pe=d // 2,
+        BLOCK_SIZE=block_size,
+        X_SIZE=x_cache if not flash_layout else 0,
+        FLASH_LAYOUT=flash_layout,
+        VALUE_SHUFFLE_LAYOUT=value_shuffle_layout,
+        HAVE_POS=(offs is not None),
+        HAVE_K_SCALE=(k_scale is not None and apply_scale),
+        HAVE_V_SCALE=(v_scale is not None and apply_scale),
+        HAVE_ZEROS=output_zeros,
+        HAS_SWA=(swa_slot_mapping is not None),
+        num_warps=1,
+    )
+
+    if zeros_out is not None:
+        return q_out.view(-1, qh * d), k_out, key_cache, value_cache, zeros_out
+    return q_out.view(-1, qh * d), k_out, key_cache, value_cache

python/sglang/kernels/ops/kvcache/trtllm_fp8_kv_kernel.py

@@ -0,0 +1,504 @@
+"""
+Fused FP8 quantization + paged KV cache write kernel for TRTLLM MHA backend.
+
+This kernel fuses the following operations:
+1. FP8 quantization of K and V tensors (from BF16/FP16 to FP8)
+2. Per-token or per-page scale computation
+3. Writing quantized K/V to paged KV cache layout
+
+Performance benefits:
+- Eliminates intermediate FP8 tensors in memory
+- Reduces kernel launch overhead
+- Better memory bandwidth utilization
+"""
+
+import logging
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _process_kv_tensor(
+    token_id,
+    head_block_id,
+    page_id,
+    page_offset,
+    input_ptr,
+    cache_ptr,
+    inv_scale,
+    use_provided_scale: tl.constexpr,
+    num_kv_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    input_stride_token: tl.constexpr,
+    input_stride_head: tl.constexpr,
+    input_stride_dim: tl.constexpr,
+    cache_stride_page: tl.constexpr,
+    cache_stride_offset: tl.constexpr,
+    cache_stride_head: tl.constexpr,
+    cache_stride_dim: tl.constexpr,
+    BLOCK_HEAD: tl.constexpr,
+    BLOCK_DIM: tl.constexpr,
+):
+    """Process a block of heads for a single K or V tensor."""
+    head_idx = head_block_id * BLOCK_HEAD
+    num_heads_in_block = min(BLOCK_HEAD, num_kv_heads - head_idx)
+
+    for dim_idx in range(0, head_dim, BLOCK_DIM):
+        num_dims_in_block = min(BLOCK_DIM, head_dim - dim_idx)
+
+        head_offsets = head_idx + tl.arange(0, BLOCK_HEAD)
+        dim_offsets = dim_idx + tl.arange(0, BLOCK_DIM)
+
+        head_mask = head_offsets < (head_idx + num_heads_in_block)
+        dim_mask = dim_offsets < (dim_idx + num_dims_in_block)
+
+        # Load from input using 3D strides
+        input_offsets = (
+            token_id * input_stride_token
+            + head_offsets[:, None] * input_stride_head
+            + dim_offsets[None, :] * input_stride_dim
+        )
+        mask = head_mask[:, None] & dim_mask[None, :]
+
+        block = tl.load(input_ptr + input_offsets, mask=mask, other=0.0)
+
+        # Quantize to FP8
+        if use_provided_scale:
+            block_fp8 = (block * inv_scale).to(tl.float8e4nv)
+        else:
+            block_fp8 = block.to(tl.float8e4nv)
+
+        # Write to cache at [page_id, page_offset, head, dim]
+        cache_offsets = (
+            page_id * cache_stride_page
+            + page_offset * cache_stride_offset
+            + head_offsets[:, None] * cache_stride_head
+            + dim_offsets[None, :] * cache_stride_dim
+        )
+
+        tl.store(cache_ptr + cache_offsets, block_fp8, mask=mask)
+
+
+@triton.jit
+def _fused_fp8_set_kv_buffer_kernel(
+    # Input tensors (post-RoPE K and V in FP16/BF16)
+    k_ptr,  # [num_tokens, num_kv_heads, head_dim]
+    v_ptr,  # [num_tokens, num_kv_heads, head_dim]
+    # Output KV cache buffers (FP8 paged layout)
+    k_cache_ptr,  # [total_slots, num_kv_heads, head_dim]
+    v_cache_ptr,  # [total_slots, num_kv_heads, head_dim]
+    # Cache location indices
+    cache_loc_ptr,  # [num_tokens] -> token to cache location mapping
+    # Pointers to scalar inverse scales (computed on GPU in wrapper)
+    inv_k_scale_ptr,  # pointer to 0-D tensor on GPU
+    inv_v_scale_ptr,  # pointer to 0-D tensor on GPU
+    use_provided_scale: tl.constexpr,  # whether to use provided scale
+    # Tensor dimensions
+    num_kv_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    page_size: tl.constexpr,
+    # Strides for K input [num_tokens, num_kv_heads, head_dim]
+    k_stride_token: tl.constexpr,
+    k_stride_head: tl.constexpr,
+    k_stride_dim: tl.constexpr,
+    # Strides for K cache [total_slots, num_kv_heads, head_dim] (logically paged)
+    k_cache_stride_page: tl.constexpr,
+    k_cache_stride_offset: tl.constexpr,
+    k_cache_stride_head: tl.constexpr,
+    k_cache_stride_dim: tl.constexpr,
+    # Strides for V input [num_tokens, num_kv_heads, head_dim]
+    v_stride_token: tl.constexpr,
+    v_stride_head: tl.constexpr,
+    v_stride_dim: tl.constexpr,
+    # Strides for V cache [total_slots, num_kv_heads, head_dim] (logically paged)
+    v_cache_stride_page: tl.constexpr,
+    v_cache_stride_offset: tl.constexpr,
+    v_cache_stride_head: tl.constexpr,
+    v_cache_stride_dim: tl.constexpr,
+    # Block sizes
+    BLOCK_HEAD: tl.constexpr,  # Number of heads per block
+    BLOCK_DIM: tl.constexpr,  # Head dimension block size
+):
+    """
+    Fused FP8 quantization + paged KV cache write kernel.
+
+    Each program processes one token-head_block-kv combination, quantizing and writing
+    to the appropriate page in the KV cache.
+
+    Grid: (num_tokens, num_head_blocks, 2) where dim2: 0=K, 1=V
+    """
+    # Get program IDs
+    token_id = tl.program_id(0)
+    head_block_id = tl.program_id(1)
+    kv_idx = tl.program_id(2)  # 0 for K, 1 for V
+
+    # Get cache location for this token
+    cache_loc = tl.load(cache_loc_ptr + token_id)
+
+    # Compute page_id and offset within page
+    page_id = cache_loc // page_size
+    page_offset = cache_loc % page_size
+
+    # Select K or V based on kv_idx
+    if kv_idx == 0:
+        # Process K tensor
+        if use_provided_scale:
+            inv_scale = tl.load(inv_k_scale_ptr)
+        else:
+            inv_scale = 1.0
+        _process_kv_tensor(
+            token_id,
+            head_block_id,
+            page_id,
+            page_offset,
+            k_ptr,
+            k_cache_ptr,
+            inv_scale,
+            use_provided_scale,
+            num_kv_heads,
+            head_dim,
+            k_stride_token,
+            k_stride_head,
+            k_stride_dim,
+            k_cache_stride_page,
+            k_cache_stride_offset,
+            k_cache_stride_head,
+            k_cache_stride_dim,
+            BLOCK_HEAD,
+            BLOCK_DIM,
+        )
+    else:
+        # Process V tensor
+        if use_provided_scale:
+            inv_scale = tl.load(inv_v_scale_ptr)
+        else:
+            inv_scale = 1.0
+        _process_kv_tensor(
+            token_id,
+            head_block_id,
+            page_id,
+            page_offset,
+            v_ptr,
+            v_cache_ptr,
+            inv_scale,
+            use_provided_scale,
+            num_kv_heads,
+            head_dim,
+            v_stride_token,
+            v_stride_head,
+            v_stride_dim,
+            v_cache_stride_page,
+            v_cache_stride_offset,
+            v_cache_stride_head,
+            v_cache_stride_dim,
+            BLOCK_HEAD,
+            BLOCK_DIM,
+        )
+
+
+def fused_fp8_set_kv_buffer(
+    k: torch.Tensor,  # [num_tokens, num_kv_heads, head_dim] or [num_tokens, num_kv_heads * head_dim]
+    v: torch.Tensor,  # [num_tokens, num_kv_heads, head_dim] or [num_tokens, num_kv_heads * head_dim]
+    k_cache: torch.Tensor,  # [total_slots, num_kv_heads, head_dim] or [num_pages, page_size, num_kv_heads, head_dim]
+    v_cache: torch.Tensor,  # [total_slots, num_kv_heads, head_dim] or [num_pages, page_size, num_kv_heads, head_dim]
+    cache_loc: torch.Tensor,  # [num_tokens], dtype=int32
+    k_scale: Optional[
+        float
+    ] = None,  # Scalar scale (matching original set_kv_buffer signature)
+    v_scale: Optional[float] = None,
+    page_size: int = 16,
+    use_triton: bool = True,  # Whether to use Triton kernel (set to False to force naive fallback)
+) -> None:
+    """
+    Python wrapper for the fused FP8 quantization + paged KV cache write kernel.
+
+    This function replicates the exact behavior of the original set_kv_buffer but with
+    a fused kernel that combines FP8 quantization and cache write.
+
+    Args:
+        k: Key tensor after RoPE, can be 2D or 3D
+        v: Value tensor, can be 2D or 3D
+        k_cache: Paged K cache buffer in FP8
+        v_cache: Paged V cache buffer in FP8
+        cache_loc: Cache location for each token, shape [num_tokens]
+        k_scale: Optional scalar scale for K (matching original set_kv_buffer)
+        v_scale: Optional scalar scale for V (matching original set_kv_buffer)
+        page_size: Number of tokens per page
+        use_triton: Whether to use optimized Triton kernel
+    """
+    num_tokens = k.shape[0]
+
+    # Step 1: Infer num_kv_heads and head_dim from cache shape
+    if k_cache.ndim == 3:
+        # 3D cache layout: [total_slots, num_kv_heads, head_dim]
+        total_slots, num_kv_heads, head_dim = k_cache.shape
+        assert (
+            total_slots % page_size == 0
+        ), f"total_slots ({total_slots}) must be divisible by page_size ({page_size})"
+        num_pages = total_slots // page_size
+    elif k_cache.ndim == 4:
+        # 4D cache layout: [num_pages, page_size, num_kv_heads, head_dim]
+        num_pages, ps, num_kv_heads, head_dim = k_cache.shape
+        assert (
+            ps == page_size
+        ), f"page_size mismatch: cache has {ps}, expected {page_size}"
+        total_slots = num_pages * page_size
+    else:
+        raise ValueError(f"Unsupported k_cache.ndim={k_cache.ndim}, expected 3 or 4")
+
+    # Step 2: Validate k, v shapes and normalize
+    # Store original 3D shape for Triton path
+    k_3d = None
+    v_3d = None
+
+    if k.ndim == 3:
+        # Input is [num_tokens, num_kv_heads, head_dim]
+        assert (
+            k.shape[1] == num_kv_heads
+        ), f"num_kv_heads mismatch: k.shape[1]={k.shape[1]} vs cache={num_kv_heads}"
+        assert (
+            k.shape[2] == head_dim
+        ), f"head_dim mismatch: k.shape[2]={k.shape[2]} vs cache={head_dim}"
+        assert v.shape[1] == num_kv_heads and v.shape[2] == head_dim, "v shape mismatch"
+
+        # Keep 3D for Triton kernel
+        k_3d = k
+        v_3d = v
+        # Create 2D view for naive fallback (will be used only if use_triton=False)
+        k_2d = k.reshape(num_tokens, num_kv_heads * head_dim)
+        v_2d = v.reshape(num_tokens, num_kv_heads * head_dim)
+    elif k.ndim == 2:
+        # Input is already [num_tokens, num_kv_heads * head_dim]
+        assert (
+            k.shape[1] == num_kv_heads * head_dim
+        ), f"k.shape[1]={k.shape[1]} != {num_kv_heads * head_dim}"
+        assert (
+            v.shape[1] == num_kv_heads * head_dim
+        ), f"v.shape[1]={v.shape[1]} != {num_kv_heads * head_dim}"
+
+        # Create 3D view for Triton kernel
+        k_3d = k.view(num_tokens, num_kv_heads, head_dim)
+        v_3d = v.view(num_tokens, num_kv_heads, head_dim)
+        # Keep 2D for naive
+        k_2d = k
+        v_2d = v
+    else:
+        raise ValueError(f"Unsupported k.ndim={k.ndim}, expected 2 or 3")
+
+    # Step 3: Compute cache strides based on layout
+    if k_cache.ndim == 3:
+        # 3D cache: [total_slots, num_kv_heads, head_dim]
+        stride_slot = k_cache.stride(0)
+        stride_head = k_cache.stride(1)
+        stride_dim = k_cache.stride(2)
+
+        k_cache_stride_page = stride_slot * page_size
+        k_cache_stride_offset = stride_slot
+        k_cache_stride_head = stride_head
+        k_cache_stride_dim = stride_dim
+
+        v_stride_slot = v_cache.stride(0)
+        v_stride_head = v_cache.stride(1)
+        v_stride_dim = v_cache.stride(2)
+
+        v_cache_stride_page = v_stride_slot * page_size
+        v_cache_stride_offset = v_stride_slot
+        v_cache_stride_head = v_stride_head
+        v_cache_stride_dim = v_stride_dim
+    else:
+        # 4D cache: [num_pages, page_size, num_kv_heads, head_dim]
+        k_cache_stride_page = k_cache.stride(0)
+        k_cache_stride_offset = k_cache.stride(1)
+        k_cache_stride_head = k_cache.stride(2)
+        k_cache_stride_dim = k_cache.stride(3)
+
+        v_cache_stride_page = v_cache.stride(0)
+        v_cache_stride_offset = v_cache.stride(1)
+        v_cache_stride_head = v_cache.stride(2)
+        v_cache_stride_dim = v_cache.stride(3)
+
+    # Decide whether to use provided scale
+    use_provided_scale = k_scale is not None and v_scale is not None
+
+    if use_triton and num_tokens > 0:
+        # Use optimized Triton kernel
+        # Compute input strides for 3D k, v: [num_tokens, num_kv_heads, head_dim]
+        k_stride_token = k_3d.stride(0)
+        k_stride_head = k_3d.stride(1)
+        k_stride_dim = k_3d.stride(2)
+
+        v_stride_token = v_3d.stride(0)
+        v_stride_head = v_3d.stride(1)
+        v_stride_dim = v_3d.stride(2)
+
+        # Block sizes for tiling (tunable)
+        BLOCK_HEAD = min(num_kv_heads, 8)  # Process up to 8 heads at once
+        BLOCK_DIM = min(head_dim, 128)  # Process up to 128 dims at once
+
+        # Compute number of head blocks
+        num_head_blocks = (num_kv_heads + BLOCK_HEAD - 1) // BLOCK_HEAD
+
+        # Grid: (num_tokens, num_head_blocks, 2)
+        # - dim 0: tokens
+        # - dim 1: head blocks
+        # - dim 2: K/V (0=K, 1=V)
+        grid = (num_tokens, num_head_blocks, 2)
+
+        device = k_3d.device
+
+        def _to_tensor_scale(scale):
+            """Convert scale to 0-D CUDA tensor (accepts Python float or Tensor)."""
+            if isinstance(scale, torch.Tensor):
+                return scale.to(device=device, dtype=torch.float32)
+            else:
+                # Python float / np scalar
+                return torch.tensor(float(scale), device=device, dtype=torch.float32)
+
+        # Compute inverse scales on GPU to avoid GPU→CPU sync in CUDA graph capture.
+        # Previously we used float(k_scale) which triggers synchronization and fails
+        # during CUDA graph capture with cudaErrorStreamCaptureUnsupported.
+        if use_provided_scale:
+            k_scale_tensor = _to_tensor_scale(k_scale)
+            v_scale_tensor = _to_tensor_scale(v_scale)
+
+            # Pure GPU scalar operation, safe for CUDA graph
+            inv_k_scale = (1.0 / k_scale_tensor).to(device=device, dtype=torch.float32)
+            inv_v_scale = (1.0 / v_scale_tensor).to(device=device, dtype=torch.float32)
+
+            inv_k_scale_ptr = inv_k_scale
+            inv_v_scale_ptr = inv_v_scale
+        else:
+            # When use_provided_scale=False, kernel uses constant 1.0 for inv_scale.
+            # Triton will optimize away the tl.load() calls via constant folding.
+            # We pass dummy pointers (k_3d) which won't be accessed in the kernel.
+            # This avoids creating new GPU tensors during CUDA graph capture.
+            inv_k_scale_ptr = k_3d
+            inv_v_scale_ptr = k_3d
+
+        # Launch Triton kernel
+        _fused_fp8_set_kv_buffer_kernel[grid](
+            k_3d,
+            v_3d,
+            k_cache,
+            v_cache,
+            cache_loc,
+            inv_k_scale_ptr,
+            inv_v_scale_ptr,
+            use_provided_scale,
+            num_kv_heads,
+            head_dim,
+            page_size,
+            k_stride_token,
+            k_stride_head,
+            k_stride_dim,
+            k_cache_stride_page,
+            k_cache_stride_offset,
+            k_cache_stride_head,
+            k_cache_stride_dim,
+            v_stride_token,
+            v_stride_head,
+            v_stride_dim,
+            v_cache_stride_page,
+            v_cache_stride_offset,
+            v_cache_stride_head,
+            v_cache_stride_dim,
+            BLOCK_HEAD=BLOCK_HEAD,
+            BLOCK_DIM=BLOCK_DIM,
+        )
+    else:
+        # Fallback to naive implementation
+        _naive_fp8_set_kv_buffer(
+            k_2d, v_2d, k_cache, v_cache, cache_loc, k_scale, v_scale, page_size
+        )
+
+
+def _naive_fp8_set_kv_buffer(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    cache_loc: torch.Tensor,
+    k_scale: Optional[float],
+    v_scale: Optional[float],
+    page_size: int,
+) -> None:
+    """
+    Naive fallback implementation that mimics the original set_kv_buffer logic.
+
+    This directly replicates the behavior of MHATokenToKVPool.set_kv_buffer:
+    1. Apply scale (if k.dtype != cache.dtype and scale is provided)
+    2. Convert to FP8
+    3. Write to cache at cache_loc
+
+    Args:
+        k: [num_tokens, num_kv_heads * head_dim], already reshaped to 2D
+        v: [num_tokens, num_kv_heads * head_dim], already reshaped to 2D
+        k_cache: [total_slots, num_kv_heads, head_dim] or [num_pages, page_size, num_kv_heads, head_dim]
+        v_cache: Same shape as k_cache
+        cache_loc: [num_tokens]
+        k_scale: Optional scale for K
+        v_scale: Optional scale for V
+        page_size: Tokens per page
+    """
+    num_tokens = k.shape[0]
+
+    # Infer dimensions from cache
+    if k_cache.ndim == 3:
+        num_kv_heads = k_cache.shape[1]
+        head_dim = k_cache.shape[2]
+    elif k_cache.ndim == 4:
+        num_kv_heads = k_cache.shape[2]
+        head_dim = k_cache.shape[3]
+    else:
+        raise ValueError(f"Unsupported k_cache.ndim={k_cache.ndim}")
+
+    # Determine target dtype and storage dtype
+    # See: python/sglang/srt/mem_cache/memory_pool.py:445-449
+    store_dtype = k_cache.dtype
+    if store_dtype == torch.uint8:
+        # Cache is stored as uint8 for FP8 (due to index_put limitation)
+        dtype = torch.float8_e4m3fn  # Logical dtype
+    else:
+        dtype = store_dtype  # Cache dtype is the logical dtype
+
+    # Replicate the original set_kv_buffer behavior
+    # See: python/sglang/srt/mem_cache/memory_pool.py:777-799
+    if k.dtype != dtype:
+        # Need quantization - clone first to avoid modifying input
+        k = k.clone()
+        v = v.clone()
+
+        if k_scale is not None:
+            k.div_(k_scale)  # In-place division
+        if v_scale is not None:
+            v.div_(v_scale)  # In-place division
+
+        k = k.to(dtype)
+        v = v.to(dtype)
+
+    # View FP8 as uint8 if needed (for index_put compatibility)
+    if store_dtype == torch.uint8 and dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+        k = k.view(torch.uint8)
+        v = v.view(torch.uint8)
+
+    # Reshape from [T, H*D] to [T, H, D]
+    k = k.view(num_tokens, num_kv_heads, head_dim)
+    v = v.view(num_tokens, num_kv_heads, head_dim)
+
+    # Write to cache using advanced indexing (same as original)
+    if k_cache.ndim == 3:
+        # 3D cache: [total_slots, H, D]
+        k_cache[cache_loc] = k
+        v_cache[cache_loc] = v
+    else:
+        # 4D cache: [num_pages, page_size, H, D]
+        # Decompose loc into page_id and page_offset (vectorized)
+        page_ids = cache_loc // page_size
+        page_offsets = cache_loc % page_size
+        k_cache[page_ids, page_offsets] = k
+        v_cache[page_ids, page_offsets] = v

python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py

@@ -0,0 +1,198 @@
+"""Fused CUDA-graph metadata update for the TRTLLM MHA backend.
+
+`TRTLLMHAAttnBackend._apply_cuda_graph_metadata` used to rebuild the
+page table(s) and seqlen buffers with ~25 small aten ops per graph
+replay (index gathers, floor_divide, cumsum, dtype casts, copies).
+On some CPUs that is ~0.7-1.0 ms of pure host dispatch, repeated 4x
+per decode step (2 draft-decode steps + target-verify + draft-extend)
+on every TP rank. The resulting per-rank CPU jitter skews the
+cudaGraphLaunch across ranks and is paid as spin time inside the first
+custom all-reduce of every replayed graph.
+
+This kernel performs the whole update in ONE launch:
+  - cache_seqlens[i]  = seq_lens[i] + seqlen_offset            (int32)
+  - cu_seqlens_k[1:]  = cumsum(cache_seqlens)                  (int32)
+  - cu_seqlens_q[1:]  = cumsum(qlens) or arange*q_stride       (optional)
+  - page_table[i, p]  = req_to_token[req_pool_indices[i],
+                                     p * page_size] // page_size
+  - swa_page_table    = full_to_swa_mapping[token] // page_size (optional)
+  - swa_out_cache_loc = full_to_swa_mapping[out_cache_loc], zero padded
+                                                               (optional)
+"""
+
+import triton
+import triton.language as tl
+
+# cu_seqlens_q handling inside the fused kernel
+Q_MODE_NONE = 0  # cu_seqlens_q is preset (decode / target-verify)
+Q_MODE_CUMSUM = 1  # cu_seqlens_q[1:] = cumsum(qlens)   (draft-extend)
+Q_MODE_STRIDED = 2  # cu_seqlens_q[1:] = arange*q_stride (draft-extend v2)
+
+
+@triton.jit
+def update_trtllm_mha_graph_metadata_kernel(
+    # inputs
+    req_pool_indices_ptr,  # [bs] int
+    seq_lens_ptr,  # [bs] int
+    req_to_token_ptr,  # [pool_size, req_to_token_stride] int32
+    swa_mapping_ptr,  # [full_size + page_size + 1] int64, or None
+    out_cache_loc_ptr,  # [num_out_tokens] int64, or None
+    qlens_ptr,  # [bs] int, or None (Q_MODE_CUMSUM only)
+    # outputs
+    cache_seqlens_ptr,  # [bs] int32
+    cu_seqlens_k_ptr,  # [bs + 1] int32
+    cu_seqlens_q_ptr,  # [bs + 1] int32, or None
+    page_table_ptr,  # [bs, page_table_stride] int32
+    swa_page_table_ptr,  # [bs, swa_page_table_stride] int32, or None
+    swa_out_cache_loc_ptr,  # [swa_out_len] int64, or None
+    # scalars
+    bs,
+    seqlen_offset,  # added to seq_lens for cache_seqlens / cu_seqlens_k
+    max_seq_pages,  # page-table columns to (re)write per row
+    q_stride,  # Q_MODE_STRIDED stride
+    num_out_tokens,  # valid prefix of out_cache_loc
+    swa_out_len,  # full swa_out_cache_loc length (zero-padded tail)
+    req_to_token_stride,
+    page_table_stride,
+    swa_page_table_stride,
+    # constexpr
+    PAGE_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+    HAS_SWA_OUT: tl.constexpr,
+    Q_MODE: tl.constexpr,
+    PAGE_BLOCK: tl.constexpr,
+    BS_BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+
+    if pid < bs:
+        # One program per batch row: cache_seqlens + page table row(s).
+        req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+        seqlen = (tl.load(seq_lens_ptr + pid) + seqlen_offset).to(tl.int32)
+        tl.store(cache_seqlens_ptr + pid, seqlen)
+
+        row_in = req_to_token_ptr + req_pool_index * req_to_token_stride
+        row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
+        if HAS_SWA:
+            swa_row_out = swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
+        for i in range(tl.cdiv(max_seq_pages, PAGE_BLOCK)):
+            page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
+            mask = page_idx < max_seq_pages
+            token = tl.load(
+                row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0
+            )
+            tl.store(row_out + page_idx, token // PAGE_SIZE, mask=mask)
+            if HAS_SWA:
+                token64 = token.to(tl.int64)
+                # Real req_to_token slots are >=0; the token>=0 guard + other=-1 mirror
+                # the swa_out_cache_loc -1 sentinel (uniform handling, no wrap).
+                swa_token = tl.load(
+                    swa_mapping_ptr + token64, mask=mask & (token64 >= 0), other=-1
+                )
+                swa_page = tl.where(swa_token < 0, -1, swa_token // PAGE_SIZE)
+                tl.store(swa_row_out + page_idx, swa_page.to(tl.int32), mask=mask)
+    elif pid == bs:
+        # Single program: cu_seqlens_k (+ optional cu_seqlens_q) cumsum.
+        offs = tl.arange(0, BS_BLOCK)
+        mask = offs < bs
+        seqlens = (tl.load(seq_lens_ptr + offs, mask=mask, other=0)).to(tl.int32)
+        seqlens = tl.where(mask, seqlens + seqlen_offset, 0)
+        tl.store(cu_seqlens_k_ptr + 1 + offs, tl.cumsum(seqlens, axis=0), mask=mask)
+        if Q_MODE == 1:  # Q_MODE_CUMSUM
+            qlens = tl.load(qlens_ptr + offs, mask=mask, other=0).to(tl.int32)
+            qlens = tl.where(mask, qlens, 0)
+            tl.store(cu_seqlens_q_ptr + 1 + offs, tl.cumsum(qlens, axis=0), mask=mask)
+        if Q_MODE == 2:  # Q_MODE_STRIDED
+            tl.store(
+                cu_seqlens_q_ptr + 1 + offs,
+                ((offs + 1) * q_stride).to(tl.int32),
+                mask=mask,
+            )
+    else:
+        # Remaining programs: swa_out_cache_loc translate + zero padding.
+        if HAS_SWA_OUT:
+            out_idx = (pid - bs - 1) * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
+            in_range = out_idx < swa_out_len
+            is_real = in_range & (out_idx < num_out_tokens)
+            loc = tl.load(out_cache_loc_ptr + out_idx, mask=is_real, other=0)
+            translated = tl.load(
+                swa_mapping_ptr + loc, mask=is_real & (loc >= 0), other=0
+            )
+            translated = tl.where(is_real & (loc < 0), -1, translated)
+            tl.store(swa_out_cache_loc_ptr + out_idx, translated, mask=in_range)
+
+
+def update_trtllm_mha_graph_metadata(
+    *,
+    req_pool_indices,
+    seq_lens,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table,
+    bs: int,
+    seqlen_offset: int,
+    max_seq_pages: int,
+    page_size: int,
+    swa_mapping=None,
+    swa_page_table=None,
+    out_cache_loc=None,
+    swa_out_cache_loc=None,
+    cu_seqlens_q=None,
+    qlens=None,
+    q_stride: int = 0,
+    q_mode: int = Q_MODE_NONE,
+):
+    """Launch the fused metadata update (one kernel for the whole replay init)."""
+    if bs == 0:
+        return
+
+    # Launch-block width: page-table columns each program writes per iteration
+    # (also the swa_out_cache_loc tile width). 512 keeps the per-program working
+    # set small enough to stay off the register-pressure / occupancy cliff while
+    # being wide enough to cover the static page-table width in few iterations.
+    PAGE_BLOCK = 512
+    has_swa = swa_page_table is not None
+    has_swa_out = swa_out_cache_loc is not None
+
+    swa_out_len = swa_out_cache_loc.shape[0] if has_swa_out else 0
+    if has_swa_out and out_cache_loc is not None:
+        num_out_tokens = min(swa_out_len, out_cache_loc.shape[0])
+    else:
+        num_out_tokens = 0
+    if num_out_tokens == 0:
+        # All loads are masked out; pass a valid dummy pointer for codegen.
+        out_cache_loc = swa_out_cache_loc
+
+    grid_extra = triton.cdiv(swa_out_len, PAGE_BLOCK) if has_swa_out else 0
+    grid = (bs + 1 + grid_extra,)
+
+    update_trtllm_mha_graph_metadata_kernel[grid](
+        req_pool_indices,
+        seq_lens,
+        req_to_token,
+        swa_mapping,
+        out_cache_loc,
+        qlens,
+        cache_seqlens,
+        cu_seqlens_k,
+        cu_seqlens_q,
+        page_table,
+        swa_page_table,
+        swa_out_cache_loc,
+        bs,
+        seqlen_offset,
+        max_seq_pages,
+        q_stride,
+        num_out_tokens,
+        swa_out_len,
+        req_to_token.stride(0),
+        page_table.stride(0),
+        swa_page_table.stride(0) if has_swa else 0,
+        PAGE_SIZE=page_size,
+        HAS_SWA=has_swa,
+        HAS_SWA_OUT=has_swa_out,
+        Q_MODE=q_mode,
+        PAGE_BLOCK=PAGE_BLOCK,
+        BS_BLOCK=triton.next_power_of_2(bs),
+    )

python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py

@@ -0,0 +1,138 @@
+"""Device-side page-table builder for the trtllm_mha attention backend.
+
+trtllm_mha builds its block (page) table from the global ``req_to_token`` pool.
+Doing it with a host-max PyTorch gather forces a ``seq_lens.max().item()`` D2H
+sync (the CPU must know the page-table width before launching). This kernel
+instead derives the per-request page count from the device-side ``seq_lens``
+tensor, so the build is sync-free: the grid/buffer use the static
+``max_num_pages`` upper bound, while each program self-guards on the real length.
+
+The kernel is MHA-owned (no dependency on the MLA kv-index kernels) and also
+emits the SWA-translated block table in the same pass via the full->SWA lookup
+table, so SWA hybrid models stay sync-free too.
+"""
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+# Tokens covered per CTA along the page-block (grid axis-1) dimension.
+# Must be a multiple of page_size (asserted in build_trtllm_mha_page_table).
+_MHA_KV_INDEX_BLOCK_TOKENS = 4096
+# Triton kernels can only read module globals that are tl.constexpr instances.
+_MHA_KV_INDEX_BLOCK_TOKENS_TL = tl.constexpr(_MHA_KV_INDEX_BLOCK_TOKENS)
+
+
+def get_num_mha_kv_index_blocks(num_pages: int, page_size: int) -> int:
+    """Grid axis-1 size: number of page-block CTAs spanning the widest sequence.
+
+    ``num_pages`` is the per-row width of the page-table buffer (the static
+    ``max_num_pages`` upper bound). One CTA handles ``_MHA_KV_INDEX_BLOCK_TOKENS
+    // page_size`` pages.
+    """
+    pages_per_block = _MHA_KV_INDEX_BLOCK_TOKENS // page_size
+    return (num_pages + pages_per_block - 1) // pages_per_block
+
+
+@triton.jit
+def create_trtllm_mha_kv_indices_triton(
+    req_to_token_ptr,  # [max_reqs, max_context_len], int32
+    req_pool_indices_ptr,  # [bs]
+    seq_lens_ptr,  # [bs], per-request KV length in tokens
+    full_to_swa_ptr,  # full->SWA token-slot lookup table, or dummy when not SWA
+    page_table_ptr,  # [bs, num_pages] int32 block ids (output)
+    swa_page_table_ptr,  # [bs, num_pages] int32 SWA block ids (output), or dummy
+    full_to_swa_numel,
+    req_to_token_stride: tl.constexpr,
+    page_table_stride: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+):
+    """Fill ``page_table_ptr`` (and ``swa_page_table_ptr`` when ``HAS_SWA``).
+
+    Program ``(pid_req, pid_blk)`` writes the block ids of request ``pid_req`` for
+    the page-block ``pid_blk``. It reads the KV token slot at each page boundary
+    from ``req_to_token`` and converts it to a block id (``slot // PAGE_SIZE``).
+    Programs past the request's page count are guarded out, so the work (and the
+    DRAM traffic) is bounded by the device-side ``seq_lens`` — no host max needed.
+
+    The SWA lookup assumes valid (``>= 0``) slots, unlike
+    ``translate_loc_from_full_to_swa``'s ``-1`` sentinel handling; page-boundary
+    reads stay within ``seq_len``, so slots are always valid here.
+    """
+    PAGES_PER_BLOCK: tl.constexpr = _MHA_KV_INDEX_BLOCK_TOKENS_TL // PAGE_SIZE
+    pid_req = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    seq_len = tl.load(seq_lens_ptr + pid_req)
+    num_pages = tl.cdiv(seq_len, PAGE_SIZE)
+    num_page_blocks = tl.cdiv(seq_len, _MHA_KV_INDEX_BLOCK_TOKENS_TL)
+    if pid_blk >= num_page_blocks:
+        return
+
+    req_pool_index = tl.load(req_pool_indices_ptr + pid_req)
+    page_idx = tl.arange(0, PAGES_PER_BLOCK) + pid_blk * PAGES_PER_BLOCK
+    token_pos = page_idx.to(tl.int64) * PAGE_SIZE
+    mask = page_idx < num_pages
+
+    slot = tl.load(
+        req_to_token_ptr
+        + req_pool_index.to(tl.int64) * req_to_token_stride
+        + token_pos,
+        mask=mask,
+    )
+    out_off = pid_req * page_table_stride + page_idx
+    tl.store(page_table_ptr + out_off, (slot // PAGE_SIZE).to(tl.int32), mask=mask)
+    if HAS_SWA:
+        swa_index = tl.minimum(tl.maximum(slot, 0), full_to_swa_numel - 1)
+        swa_slot = tl.load(full_to_swa_ptr + swa_index.to(tl.int64), mask=mask)
+        tl.store(
+            swa_page_table_ptr + out_off,
+            (swa_slot // PAGE_SIZE).to(tl.int32),
+            mask=mask,
+        )
+
+
+def build_trtllm_mha_page_table(
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    page_table: torch.Tensor,
+    page_size: int,
+    swa_page_table: Optional[torch.Tensor] = None,
+    full_to_swa: Optional[torch.Tensor] = None,
+) -> None:
+    """Fill ``page_table`` (and ``swa_page_table`` when SWA) on-device, no D2H sync.
+
+    Computes the launch grid from the static page-table width and dispatches
+    ``create_trtllm_mha_kv_indices_triton``. ``page_table`` (and, for SWA models,
+    ``swa_page_table``) are written in place; the caller owns the buffers so the
+    cuda-graph path can reuse its pre-allocated tensors. SWA is enabled iff
+    ``full_to_swa`` is provided, which then also requires ``swa_page_table``.
+    """
+    has_swa = full_to_swa is not None
+    assert has_swa == (
+        swa_page_table is not None
+    ), "full_to_swa and swa_page_table must be provided together"
+    assert (
+        _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0
+    ), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}"
+    bs, num_pages = page_table.shape
+    full_to_swa_numel = full_to_swa.numel() if has_swa else 0
+    create_trtllm_mha_kv_indices_triton[
+        (bs, get_num_mha_kv_index_blocks(num_pages, page_size))
+    ](
+        req_to_token,
+        req_pool_indices,
+        cache_seqlens,
+        full_to_swa,
+        page_table,
+        swa_page_table,
+        full_to_swa_numel,
+        req_to_token.stride(0),
+        page_table.stride(0),
+        PAGE_SIZE=page_size,
+        HAS_SWA=has_swa,
+    )

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",
+]

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"]

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.<module>``); 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__ = []

python/sglang/kernels/ops/memory/allocator.py

@@ -0,0 +1,135 @@
+import triton
+import triton.language as tl
+
+
+# free_page_ptr aliases self.free_pages, which the paged allocator re-slices
+# after every allocation (self.free_pages = self.free_pages[num_new_pages:]).
+# Slicing only advances data_ptr() by num_new_pages * 8 bytes, so the pointer
+# flips between 16-byte-aligned and unaligned across calls. Triton specializes
+# on pointer alignment by default and bakes it into the cache key, compiling two
+# kernel variants (one with tt.divisibility=16 on free_page_ptr, one without)
+# so the second prefill on a fresh DCP server hits the alternate alignment and
+# pays an extra ~100ms JIT for that kernel variant. do_not_specialize skips
+# that specialization so only one kernel is ever compiled; the perf cost is
+# negligible (this kernel runs in ~10us and only loads ~4KB through this ptr).
+@triton.jit(do_not_specialize=["free_page_ptr"])
+def alloc_extend_kernel(
+    pre_lens_ptr,
+    seq_lens_ptr,
+    last_loc_ptr,
+    free_page_ptr,
+    out_indices,
+    bs_upper: tl.constexpr,
+    page_size: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    load_offset = tl.arange(0, bs_upper)
+    seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid)
+    pre_lens = tl.load(pre_lens_ptr + load_offset, mask=load_offset <= pid)
+    extend_lens = seq_lens - pre_lens
+
+    seq_len = tl.load(seq_lens_ptr + pid)
+    pre_len = tl.load(pre_lens_ptr + pid)
+    extend_len = seq_len - pre_len
+
+    sum_extend_lens = tl.sum(extend_lens)
+    output_start_loc = sum_extend_lens - extend_len
+
+    num_pages_after = (seq_lens + page_size - 1) // page_size
+    num_pages_before = (pre_lens + page_size - 1) // page_size
+    num_new_pages = num_pages_after - num_pages_before
+
+    num_page_start_loc_self = (seq_len + page_size - 1) // page_size - (
+        pre_len + page_size - 1
+    ) // page_size
+    sum_num_new_pages = tl.sum(num_new_pages)
+    new_page_start_loc = sum_num_new_pages - num_page_start_loc_self
+
+    # Part 1: fill the old partial page
+    last_loc = tl.load(last_loc_ptr + pid)
+    num_part1 = (
+        min(seq_len, (pre_len + page_size - 1) // page_size * page_size) - pre_len
+    )
+    offset_one_page = tl.arange(0, page_size)
+    tl.store(
+        out_indices + output_start_loc + offset_one_page,
+        last_loc + 1 + offset_one_page,
+        mask=offset_one_page < num_part1,
+    )
+    if pre_len + num_part1 == seq_len:
+        return
+
+    # Part 2: fill the new full pages using a dynamic blocked loop.
+    # The loop bound is derived from num_part2 (runtime value), so Triton
+    # generates a real loop instead of unrolling -- no constexpr dependency
+    # on extend size and only one kernel compilation.
+    num_part2 = (
+        seq_len // page_size * page_size
+        - (pre_len + page_size - 1) // page_size * page_size
+    )
+    BLOCK_EXTEND: tl.constexpr = 4096
+    num_blocks = (num_part2 + BLOCK_EXTEND - 1) // BLOCK_EXTEND
+    for block_id in range(num_blocks):
+        offset_in_block = tl.arange(0, BLOCK_EXTEND)
+        offset = block_id * BLOCK_EXTEND + offset_in_block
+        mask = offset < num_part2
+        page_start = tl.load(
+            free_page_ptr + new_page_start_loc + offset // page_size,
+            mask=mask,
+        )
+        tl.store(
+            out_indices + output_start_loc + num_part1 + offset,
+            page_start * page_size + offset % page_size,
+            mask=mask,
+        )
+    if pre_len + num_part1 + num_part2 == seq_len:
+        return
+
+    # Part 3: fill the new partial page
+    num_part3 = seq_len - seq_len // page_size * page_size
+    start_loc = tl.load(
+        free_page_ptr + new_page_start_loc + num_page_start_loc_self - 1
+    )
+    tl.store(
+        out_indices + output_start_loc + num_part1 + num_part2 + offset_one_page,
+        start_loc * page_size + offset_one_page,
+        mask=offset_one_page < num_part3,
+    )
+
+
+# Same free_page_ptr alignment rationale as alloc_extend_kernel above.
+@triton.jit(do_not_specialize=["free_page_ptr"])
+def alloc_decode_kernel(
+    seq_lens_ptr,
+    last_loc_ptr,
+    free_page_ptr,
+    out_indices,
+    bs_upper: tl.constexpr,
+    page_size: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    load_offset = tl.arange(0, bs_upper)
+    seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid)
+    pre_lens = tl.where(load_offset <= pid, seq_lens - 1, seq_lens)
+
+    seq_len = tl.load(seq_lens_ptr + pid)
+    pre_len = seq_len - 1
+
+    num_pages_after = (seq_lens + page_size - 1) // page_size
+    num_pages_before = (pre_lens + page_size - 1) // page_size
+    num_new_pages = num_pages_after - num_pages_before
+
+    num_page_start_loc_self = (seq_len + page_size - 1) // page_size - (
+        pre_len + page_size - 1
+    ) // page_size
+    sum_num_new_pages = tl.sum(num_new_pages)
+    new_page_start_loc = sum_num_new_pages - num_page_start_loc_self
+
+    if num_page_start_loc_self == 0:
+        last_loc = tl.load(last_loc_ptr + pid)
+        tl.store(out_indices + pid, last_loc + 1)
+    else:
+        page = tl.load(free_page_ptr + new_page_start_loc)
+        tl.store(out_indices + pid, page * page_size)

python/sglang/kernels/ops/memory/common.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def write_req_to_token_pool_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices,
+    prefix_tensors,
+    pre_lens,
+    seq_lens,
+    extend_lens,
+    out_cache_loc,
+    req_to_token_ptr_stride: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(0)
+
+    req_pool_index = tl.load(req_pool_indices + pid)
+    pre_len = tl.load(pre_lens + pid)
+    seq_len = tl.load(seq_lens + pid)
+    prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64))
+
+    # write prefix
+    num_loop = tl.cdiv(pre_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < pre_len
+        value = tl.load(prefix_tensor + offset, mask=mask)
+        tl.store(
+            req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset,
+            value,
+            mask=mask,
+        )
+
+    # NOTE: This can be slow for large bs
+    cumsum_start = tl.cast(0, tl.int64)
+    for i in range(pid):
+        cumsum_start += tl.load(extend_lens + i)
+
+    num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < (seq_len - pre_len)
+        value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask)
+        tl.store(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + offset
+            + pre_len,
+            value,
+            mask=mask,
+        )
+
+
+@triton.jit
+def _get_last_loc_safe_kernel(
+    req_to_token,
+    req_pool_indices_tensor,
+    prefix_lens_tensor,
+    result_i32,
+    num_tokens,
+    req_to_token_stride,
+    BLOCK_SIZE: tl.constexpr,
+    PREFIX_DTYPE_IS_I64: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
+    mask = offset < num_tokens
+
+    if PREFIX_DTYPE_IS_I64:
+        prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
+        req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)
+        token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1)
+    else:
+        prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
+        req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)
+        token_index = req_pool_indices.to(tl.int64) * req_to_token_stride + (
+            prefix_lens.to(tl.int64) - 1
+        )
+
+    token_mask = mask & (prefix_lens > 0)
+    tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1)
+    # Result stays int32 (req_to_token dtype); caller promotes after return.
+    tl.store(result_i32 + offset, tokens, mask=mask)
+
+
+def get_last_loc_triton_safe(
+    req_to_token: torch.Tensor,
+    req_pool_indices_tensor: torch.Tensor,
+    prefix_lens_tensor: torch.Tensor,
+) -> torch.Tensor:
+    """Fused `last_loc` Triton kernel whose in-kernel result buffer is int32
+    (the dtype of req_to_token). The consumer-dtype promotion happens in
+    torch after the kernel returns, so Triton never issues a mixed-width
+    store -- avoiding the HIP int32->int64 store bug hit by the legacy kernel.
+    """
+    num_tokens = prefix_lens_tensor.shape[0]
+    BLOCK_SIZE = 256
+    result_i32 = torch.empty(
+        num_tokens, dtype=torch.int32, device=prefix_lens_tensor.device
+    )
+    grid = (triton.cdiv(num_tokens, BLOCK_SIZE),)
+    _get_last_loc_safe_kernel[grid](
+        req_to_token,
+        req_pool_indices_tensor,
+        prefix_lens_tensor,
+        result_i32,
+        num_tokens,
+        req_to_token.stride(0),
+        BLOCK_SIZE=BLOCK_SIZE,
+        PREFIX_DTYPE_IS_I64=(prefix_lens_tensor.dtype == torch.int64),
+    )
+    return result_i32.to(prefix_lens_tensor.dtype)
+
+
+@triton.jit
+def get_last_loc_kernel(
+    req_to_token,
+    req_pool_indices_tensor,
+    prefix_lens_tensor,
+    result,
+    num_tokens,
+    req_to_token_stride,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
+    mask = offset < num_tokens
+
+    prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
+    req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)
+
+    token_mask = prefix_lens > 0
+    token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1)
+    tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1)
+
+    tl.store(result + offset, tokens, mask=mask)
+
+
+def get_last_loc_triton(
+    req_to_token: torch.Tensor,
+    req_pool_indices_tensor: torch.Tensor,
+    prefix_lens_tensor: torch.Tensor,
+) -> torch.Tensor:
+    BLOCK_SIZE = 256
+    num_tokens = prefix_lens_tensor.shape[0]
+    result = torch.empty_like(prefix_lens_tensor)
+    grid = (triton.cdiv(num_tokens, BLOCK_SIZE),)
+
+    get_last_loc_kernel[grid](
+        req_to_token,
+        req_pool_indices_tensor,
+        prefix_lens_tensor,
+        result,
+        num_tokens,
+        req_to_token.stride(0),
+        BLOCK_SIZE,
+    )
+    return result

python/sglang/kernels/ops/memory/virtual_slot.py

@@ -0,0 +1,96 @@
+# Copyright 2023-2026 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.
+# ==============================================================================
+"""Virtual<->physical slot Triton kernels for the unified memory pool."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+# Fused take-physical-pages + bind for the alloc fast path. Invoked ONLY when
+# `_hole_count == 0`; otherwise the slow path drains holes first (Invariant B,
+# greedy hole reuse). Caller advances `watermark_physical` and checks overflow
+# BEFORE launch, passing the PRE-extension watermark. Cuda-graph safe (no
+# `.item()`, no tensor branching); runs on the scheduler thread.
+
+
+@triton.jit
+def alloc_bind_inplace_kernel(
+    v_pages_ptr,  # in: [N] int64 — virtual page ids
+    v2p_ptr,  # in/out: int64 — virtual_to_physical table
+    p2v_ptr,  # in/out: int64 — physical_to_virtual table
+    out_phys_ptr,  # out: [N] int64 — physical page ids
+    N,  # runtime: number of pages to allocate
+    start_phys,  # runtime: lowest physical page id in the new range
+    BLOCK: tl.constexpr,
+):
+    """Fused: ascending arange + out_phys/v2p/p2v scatter.
+
+    Caller pre-adjusts `start_phys` per direction so the range is always
+    ascending (grow-up: start_wm; grow-down: start_wm - N + 1), making the
+    v->p mapping byte-identical to the `torch.arange` slow path.
+    """
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < N
+
+    phys = (start_phys + offs).to(tl.int64)
+    v = tl.load(v_pages_ptr + offs, mask=mask, other=0).to(tl.int64)
+
+    # Masked stores skip out-of-range lanes, and `other=0` keeps us off the
+    # v2p[0]/p2v[0] padding-sink slot.
+    tl.store(out_phys_ptr + offs, phys, mask=mask)
+    tl.store(v2p_ptr + v, phys, mask=mask)
+    tl.store(p2v_ptr + phys, v, mask=mask)
+
+
+ALLOC_BIND_BLOCK = 128
+
+
+def alloc_bind_inplace(
+    v_pages: torch.Tensor,
+    v2p: torch.Tensor,
+    p2v: torch.Tensor,
+    start_phys: int,
+) -> torch.Tensor:
+    """Allocate N ascending physical pages from `start_phys` and bind to `v_pages`.
+
+    Caller must advance `watermark_physical` by N and verify overflow BEFORE
+    calling; this launcher does neither.
+    """
+    N = int(v_pages.numel())
+    if N == 0:
+        return torch.empty(0, dtype=torch.int64, device=v_pages.device)
+    if not v_pages.is_cuda:
+        # Pure-torch CPU reference for the CUDA-only kernel.
+        phys_pages = torch.arange(
+            start_phys, start_phys + N, dtype=torch.int64, device=v_pages.device
+        )
+        v = v_pages.to(torch.int64)
+        v2p[v] = phys_pages
+        p2v[phys_pages] = v
+        return phys_pages
+    phys_pages = torch.empty(N, dtype=torch.int64, device=v_pages.device)
+    grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),)
+    alloc_bind_inplace_kernel[grid](
+        v_pages,
+        v2p,
+        p2v,
+        phys_pages,
+        N,
+        start_phys,
+        BLOCK=ALLOC_BIND_BLOCK,
+    )
+    return phys_pages

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

python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py

@@ -0,0 +1,701 @@
+# Temporarily adapted from https://github.com/vllm-project/vllm/blob/main/vllm/lora/ops/triton_ops/fused_moe_lora_op.py, will optimize in future refactor
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.distributed import (
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.srt.utils.common import is_blackwell_supported, is_sm90_supported
+
+# Import SGLang's standard PDL support detection
+
+
+_LORA_PTR_DICT: dict[tuple[int, ...], torch.Tensor] = {}
+
+
+def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
+    """
+    `_LORA_PTR_DICT` collects the required information during `profile_run`,
+    After this, it remains constant and subsequent usage is through LUT.
+    Refer to:
+    https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
+    """
+    key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights)
+
+    if (ptr_tensor := _LORA_PTR_DICT.get(key)) is not None:
+        return ptr_tensor
+
+    tensor_ptrs = []
+    for lora_weight in lora_weights:
+        tensor_ptrs.append(lora_weight.data_ptr())
+    ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)
+
+    _LORA_PTR_DICT[key] = ptr_tensor
+    return _LORA_PTR_DICT.get(key)
+
+
+@triton.jit(
+    do_not_specialize=[
+        "num_valid_tokens",
+        "EM",
+        "stride_tl",
+        "stride_el",
+        "slice_a_size",
+        "slice_c_size",
+    ]
+)
+def _fused_moe_lora_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    num_experts,
+    lora_ids,
+    adapter_enabled,
+    # The stride variables represent how much to increase the ptr by when
+    # moving by 1 element in a particular dimension. E.g. `stride_am` is
+    # how much to increase `a_ptr` by to get the element one row down
+    # (A has M rows).
+    stride_am,
+    stride_ak,
+    stride_bl,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_tl,
+    stride_el,
+    slice_a_size,
+    slice_c_size,
+    # Meta-parameters
+    num_slice_a: tl.constexpr,
+    num_slice_c: tl.constexpr,
+    top_k: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    USE_GDC: tl.constexpr,
+    launch_pdl: tl.constexpr,
+    IS_PRIMARY: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    slice_id = tl.program_id(axis=1)
+    lora_idx = tl.program_id(axis=2)
+    lora_id = tl.load(lora_ids + lora_idx)
+
+    if lora_id == -1:
+        # Early exit for the no-lora case.
+        return
+    moe_enabled = tl.load(adapter_enabled + lora_id)
+    if moe_enabled == 0:
+        # Early exit for the no moe lora case.
+        return
+    max_loras = tl.num_programs(axis=2)
+    grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
+
+    # calculate pid_m,pid_n
+    pid_sk = pid % SPLIT_K
+    pid_m_n = pid // SPLIT_K
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid_m_n // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
+    pid_n = (pid_m_n % num_pid_in_group) // group_size_m
+
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_id)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+    # get the expert_id to process curr shard
+    ind = lora_id * stride_el + pid_m
+    expert_id = tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)
+    if expert_id == -1:
+        return
+
+    # get a_ptr,b_ptr,c_ptr
+    cur_a_ptr = a_ptr + (slice_id % num_slice_a) * slice_a_size
+    cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(c_ptr.dtype.element_ty))
+    cur_c_ptr = c_ptr + (slice_id % num_slice_c) * slice_c_size
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+    # ================================================================= secure
+
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    token_ind = stride_tl * lora_id + offs_token_id
+    offs_token = tl.load(
+        sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0
+    )
+    token_mask = offs_token < num_valid_tokens
+
+    # ================================================================= secure
+
+    # get a_ptrs,b_ptrs
+    a_ptrs = cur_a_ptr + (
+        offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
+    )
+
+    b_ptrs = (
+        cur_b_ptr
+        + lora_id * stride_bl
+        + expert_id * stride_be
+        + offs_k[:, None] * stride_bk
+        + offs_bn[None, :] * stride_bn
+    )
+
+    if USE_GDC and IS_PRIMARY:
+        # GDC launch dependents hints the runtime system to launch dependent kernels.
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # ================================================================= secure
+
+    # accumulator
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    # ================================================================= secure
+
+    # GDC wait waits for ALL programs in the prior kernel to complete
+    # before continuing.
+    if USE_GDC and not IS_PRIMARY:
+        tl.extra.cuda.gdc_wait()
+
+    for k in range(0, grid_k):
+        k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
+        # pre-fetch lora weight
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
+        a = tl.load(
+            a_ptrs,
+            mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
+            other=0.0,
+        )
+        accumulator += tl.dot(a, b.to(a.dtype))
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
+
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+        accumulator = accumulator * moe_weight[:, None]
+    accumulator = accumulator.to(c_ptr.dtype.element_ty)
+    # Write back the block of the output
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+
+    if SPLIT_K == 1:
+        tl.store(c_ptrs, accumulator, mask=c_mask)
+    else:
+        tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")
+
+
+@torch.inference_mode()
+def _fused_moe_lora_shrink(
+    a_intermediate_cache1: torch.Tensor,
+    # (num_slices, num_tokens, top_k_num, max_lora_rank)
+    qcurr_hidden_states: torch.Tensor,  # (num_tokens, K,)
+    lora_a_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, max_lora_rank, K,),...]
+    topk_weights: torch.Tensor,  # (num_tokens, top_k_num)
+    sorted_token_ids: torch.Tensor,  # (max_loras, _)
+    expert_ids: torch.Tensor,  # (max_loras, _ ,)
+    num_tokens_post_padded: torch.Tensor,  # (max_loras, )
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    ## adding for kernel
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    top_k_divisor: int = None,
+    mul_routed_weight: bool = False,
+) -> None:
+    w1_lora_a_stacked = lora_a_stacked[0]
+
+    use_gdc = is_sm90_supported() or is_blackwell_supported()
+    shrink_config = {
+        "BLOCK_SIZE_M": block_size_m,
+        "BLOCK_SIZE_N": block_size_n,
+        "BLOCK_SIZE_K": block_size_k,
+        "GROUP_SIZE_M": group_size_m,
+        "num_warps": num_warps,
+        "num_stages": num_stages,
+        "SPLIT_K": split_k,
+        "USE_GDC": use_gdc,
+        "launch_pdl": use_gdc,  # triton kernel metadata
+    }
+
+    b_ptr = _get_ptr(lora_a_stacked, device)
+
+    grid = lambda META: (
+        split_k
+        * triton.cdiv(EM, META["BLOCK_SIZE_M"])
+        * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        len(lora_a_stacked),
+        lora_a_stacked[0].shape[0],
+    )
+    _fused_moe_lora_kernel[grid](
+        qcurr_hidden_states,
+        b_ptr,
+        a_intermediate_cache1,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        EM,
+        num_tokens,
+        num_experts,
+        lora_ids,
+        adapter_enabled,
+        qcurr_hidden_states.stride(0),
+        qcurr_hidden_states.stride(1),
+        w1_lora_a_stacked.stride(0),
+        w1_lora_a_stacked.stride(1),
+        w1_lora_a_stacked.stride(3),
+        w1_lora_a_stacked.stride(2),
+        a_intermediate_cache1.stride(2),
+        a_intermediate_cache1.stride(3),
+        sorted_token_ids.stride(0),
+        expert_ids.stride(0),
+        slice_a_size=qcurr_hidden_states.numel(),
+        slice_c_size=a_intermediate_cache1.numel() // num_slices,
+        num_slice_a=1,
+        num_slice_c=num_slices,
+        top_k=(
+            top_k_divisor
+            if top_k_divisor is not None
+            else (1 if mul_routed_weight else top_k_num)
+        ),
+        MUL_ROUTED_WEIGHT=False,
+        IS_PRIMARY=True,
+        **shrink_config,
+    )
+
+
+@torch.inference_mode()
+def _fused_moe_lora_expand(
+    output: torch.Tensor,  # (num_tokens, top_k_num, N*len(lora_a_stacked),)
+    a_intermediate_cache1: torch.Tensor,  # (num_slices, M, top_k_num, max_lora_rank)
+    b_intermediate_cache1: torch.Tensor,  # (num_slices, M, top_k_num, output_dim_size)
+    lora_b_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, max_lora_rank, K,),...]
+    topk_weights: torch.Tensor,  # (num_tokens, top_k_num)
+    sorted_token_ids: torch.Tensor,  # (max_loras, _)
+    expert_ids: torch.Tensor,  # (max_loras, _ ,)
+    num_tokens_post_padded: torch.Tensor,  # (max_loras, )
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    ## adding for kernel
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    max_lora_rank: int,
+    w1_output_dim_size: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    mul_routed_weight: bool = False,
+    offset: int = 0,
+) -> None:
+
+    b_ptr = _get_ptr(lora_b_stacked, device)
+    K = max_lora_rank
+    N = w1_output_dim_size
+
+    w1_lora_b_stacked = lora_b_stacked[0]
+
+    a_intermediate_cache1 = a_intermediate_cache1.view(
+        -1, a_intermediate_cache1.shape[3]
+    )
+
+    use_gdc = is_sm90_supported() or is_blackwell_supported()
+    expand_config = {
+        "BLOCK_SIZE_M": block_size_m,
+        "BLOCK_SIZE_N": block_size_n,
+        "BLOCK_SIZE_K": block_size_k,
+        "GROUP_SIZE_M": group_size_m,
+        "num_warps": num_warps,
+        "num_stages": num_stages,
+        "SPLIT_K": split_k,  # Set split_k = 1 for expand calls
+        "USE_GDC": use_gdc,
+        "launch_pdl": use_gdc,  # triton kernel metadata
+    }
+
+    grid = lambda META: (
+        triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        len(lora_b_stacked),
+        lora_b_stacked[0].shape[0],
+    )
+    _fused_moe_lora_kernel[grid](
+        a_intermediate_cache1,
+        b_ptr,
+        b_intermediate_cache1,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        EM,
+        num_tokens,
+        num_experts,
+        lora_ids,
+        adapter_enabled,
+        a_intermediate_cache1.stride(0),
+        a_intermediate_cache1.stride(1),
+        w1_lora_b_stacked.stride(0),
+        w1_lora_b_stacked.stride(1),
+        w1_lora_b_stacked.stride(3),
+        w1_lora_b_stacked.stride(2),
+        b_intermediate_cache1.stride(2),
+        b_intermediate_cache1.stride(3),
+        sorted_token_ids.stride(0),
+        expert_ids.stride(0),
+        slice_a_size=a_intermediate_cache1.numel() // num_slices,
+        slice_c_size=b_intermediate_cache1.numel() // num_slices,
+        num_slice_a=num_slices,
+        num_slice_c=num_slices,
+        top_k=1,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        IS_PRIMARY=False,
+        **expand_config,
+    )
+    for i in range(num_slices):
+        output[:, :, i * N + offset : (i + 1) * N + offset] += b_intermediate_cache1[i]
+
+
+@torch.inference_mode()
+def _fused_moe_lora(
+    output: torch.Tensor,  # (num_tokens, top_k_num, N*len(lora_a_stacked),)
+    qcurr_hidden_states: torch.Tensor,  # (num_tokens, K,)
+    lora_a_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, max_lora_rank, K,),...]
+    lora_b_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, N, max_lora_rank,),...]
+    topk_weights: torch.Tensor,  # (num_tokens, top_k_num)
+    sorted_token_ids: torch.Tensor,  # (max_loras, _)
+    expert_ids: torch.Tensor,  # (max_loras, _ ,)
+    num_tokens_post_padded: torch.Tensor,  # (max_loras, )
+    max_lora_rank: int,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    shrink_block_size_m: int,
+    shrink_block_size_n: int,
+    shrink_block_size_k: int,
+    shrink_group_size_m: int,
+    shrink_num_warps: int,
+    shrink_num_stages: int,
+    shrink_split_k: int,
+    expand_block_size_m: int,
+    expand_block_size_n: int,
+    expand_block_size_k: int,
+    expand_group_size_m: int,
+    expand_num_warps: int,
+    expand_num_stages: int,
+    expand_split_k: int,
+    mul_routed_weight: bool = False,
+    fully_sharded: bool = False,
+    offset: int = 0,
+) -> None:
+    assert len(lora_a_stacked) == len(lora_b_stacked) > 0
+    assert (
+        sorted_token_ids.dim()
+        == expert_ids.dim()
+        == topk_weights.dim()
+        == qcurr_hidden_states.dim()
+        == 2
+    )
+    assert (
+        sorted_token_ids.shape[0]
+        == expert_ids.shape[0]
+        == num_tokens_post_padded.shape[0]
+    )
+    assert output.shape[0] == topk_weights.shape[0]
+    assert top_k_num == topk_weights.shape[1]
+    device = qcurr_hidden_states.device
+    num_slices = len(lora_a_stacked)
+    w1_lora_b_stacked = lora_b_stacked[0]
+    num_experts = lora_a_stacked[0].shape[1]
+    N = max_lora_rank
+    M = topk_weights.shape[0]
+    EM = sorted_token_ids.shape[1]
+    K = qcurr_hidden_states.shape[1]
+    num_tokens = M * top_k_num
+    w1_output_dim_size = w1_lora_b_stacked.shape[2]
+
+    # Detect whether input is already expanded (down path: [M*top_k, dim])
+    # or not (gate_up path: [M, dim]). Down path needs divisor=1.
+    input_is_expanded = qcurr_hidden_states.shape[0] == M * top_k_num
+    shrink_top_k_divisor = 1 if input_is_expanded else top_k_num
+
+    a_intermediate_cache1 = torch.zeros(
+        (num_slices, M, top_k_num, max_lora_rank),
+        dtype=output.dtype,
+        device=device,
+    )
+
+    b_intermediate_cache1 = torch.zeros(
+        (num_slices, M, top_k_num, w1_output_dim_size),
+        dtype=output.dtype,
+        device=device,
+    )
+
+    _fused_moe_lora_shrink(
+        a_intermediate_cache1,
+        qcurr_hidden_states,
+        lora_a_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        top_k_num,
+        lora_ids,
+        adapter_enabled,
+        ## adding for kernel
+        device,
+        N,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        shrink_block_size_m,
+        shrink_block_size_n,
+        shrink_block_size_k,
+        shrink_group_size_m,
+        shrink_num_warps,
+        shrink_num_stages,
+        shrink_split_k,
+        top_k_divisor=shrink_top_k_divisor,
+        mul_routed_weight=False,
+    )
+
+    if fully_sharded:
+        if max_lora_rank == w1_lora_b_stacked.shape[-1]:
+            a_intermediate_cache1 = tensor_model_parallel_all_reduce(
+                a_intermediate_cache1
+            )
+        else:
+            a_intermediate_cache1 = tensor_model_parallel_all_gather(
+                a_intermediate_cache1
+            )
+
+            # reset max_lora_rank to the full rank after allgather
+            max_lora_rank = a_intermediate_cache1.shape[-1]
+
+    _fused_moe_lora_expand(
+        output,
+        a_intermediate_cache1,
+        b_intermediate_cache1,
+        lora_b_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        top_k_num,
+        lora_ids,
+        adapter_enabled,
+        ## adding for kernel
+        device,
+        N,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        max_lora_rank,
+        w1_output_dim_size,
+        expand_block_size_m,
+        expand_block_size_n,
+        expand_block_size_k,
+        expand_group_size_m,
+        expand_num_warps,
+        expand_num_stages,
+        expand_split_k,
+        mul_routed_weight,
+        offset,
+    )
+
+
+def _fused_moe_lora_fake(
+    output: torch.Tensor,
+    qcurr_hidden_states: torch.Tensor,
+    lora_a_stacked: list[torch.Tensor],
+    lora_b_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    max_lora_rank: int,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    shrink_block_size_m: int,
+    shrink_block_size_n: int,
+    shrink_block_size_k: int,
+    shrink_group_size_m: int,
+    shrink_num_warps: int,
+    shrink_num_stages: int,
+    shrink_split_k: int,
+    expand_block_size_m: int,
+    expand_block_size_n: int,
+    expand_block_size_k: int,
+    expand_group_size_m: int,
+    expand_num_warps: int,
+    expand_num_stages: int,
+    expand_split_k: int,
+    mul_routed_weight: bool = False,
+    fully_sharded: bool = False,
+    offset: int = 0,
+) -> None:
+    return
+
+
+def _fused_moe_lora_shrink_fake(
+    a_intermediate_cache1: torch.Tensor,
+    qcurr_hidden_states: torch.Tensor,
+    lora_a_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    mul_routed_weight: bool = False,
+) -> None:
+    return
+
+
+def _fused_moe_lora_expand_fake(
+    output: torch.Tensor,
+    a_intermediate_cache1: torch.Tensor,
+    b_intermediate_cache1: torch.Tensor,
+    lora_b_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    max_lora_rank: int,
+    w1_output_dim_size: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    mul_routed_weight: bool = False,
+    offset: int = 0,
+) -> None:
+    return
+
+
+# Register as SGLang custom ops following the same pattern as other ops
+try:
+    from sglang.srt.utils.common import direct_register_custom_op
+
+    direct_register_custom_op(
+        op_name="fused_moe_lora",
+        op_func=_fused_moe_lora,
+        mutates_args=["output"],
+        fake_impl=_fused_moe_lora_fake,
+    )
+
+    direct_register_custom_op(
+        op_name="fused_moe_lora_shrink",
+        op_func=_fused_moe_lora_shrink,
+        mutates_args=["a_intermediate_cache1"],
+        fake_impl=_fused_moe_lora_shrink_fake,
+    )
+
+    direct_register_custom_op(
+        op_name="fused_moe_lora_expand",
+        op_func=_fused_moe_lora_expand,
+        mutates_args=["output", "b_intermediate_cache1"],
+        fake_impl=_fused_moe_lora_expand_fake,
+    )
+
+    # Export through torch.ops.sglang namespace
+    fused_moe_lora = torch.ops.sglang.fused_moe_lora
+    fused_moe_lora_shrink = torch.ops.sglang.fused_moe_lora_shrink
+    fused_moe_lora_expand = torch.ops.sglang.fused_moe_lora_expand
+
+except AttributeError:
+    fused_moe_lora = _fused_moe_lora
+    fused_moe_lora_shrink = _fused_moe_lora_shrink
+    fused_moe_lora_expand = _fused_moe_lora_expand

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)."""

python/sglang/kernels/ops/moe/virtual_experts.py

@@ -0,0 +1,786 @@
+"""
+LoRA Virtual Experts Triton Ops.
+"""
+
+import functools
+from typing import Any
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.moe_align import moe_align_block_size as jit_moe_align_block_size
+
+
+@triton.jit
+def _fused_virtual_topk_ids_kernel(
+    topk_ids_ptr,
+    token_lora_mapping_ptr,
+    virtual_topk_ids_ptr,
+    token_lora_mask_ptr,
+    num_experts_for_weight: tl.constexpr,
+    M,
+    top_k: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Fuses _get_virtual_topk_ids: comparison + clamp + arithmetic into one kernel.
+
+    For each (m, k):
+        lora_id = token_lora_mapping[m]
+        mask[m] = (lora_id >= 0)
+        safe_lora = max(lora_id, 0)
+        if shared_outer:  (handled by num_experts_for_weight == 0 sentinel)
+            virtual_topk_ids[m, k] = safe_lora * 1  (= safe_lora)
+        else:
+            virtual_topk_ids[m, k] = topk_ids[m, k] + safe_lora * num_experts_for_weight
+    """
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    total = M * top_k
+    valid = offs < total
+
+    m = offs // top_k
+    # k = offs % top_k  # not needed directly
+
+    lora_id = tl.load(token_lora_mapping_ptr + m, mask=valid, other=0)
+    mask_val = lora_id >= 0
+    safe_lora = tl.maximum(lora_id, 0)
+
+    base = tl.load(topk_ids_ptr + offs, mask=valid, other=0)
+    # Preserve negative sentinel topk_ids (e.g. -1 for non-local experts after
+    # EP dispatch). Without this, `-1 + safe_lora * num_experts` would land on
+    # a real virtual-expert slot belonging to another adapter and trigger OOB
+    # loads in downstream LoRA kernels.
+    shifted = base + safe_lora * num_experts_for_weight
+    result = tl.where(base < 0, base, shifted)
+    tl.store(virtual_topk_ids_ptr + offs, result, mask=valid)
+
+    # Write mask once per row (at first k position)
+    k = offs % top_k
+    is_first_k = k == 0
+    tl.store(token_lora_mask_ptr + m, mask_val, mask=valid & is_first_k)
+
+
+def _fused_virtual_topk_ids(
+    topk_ids: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    num_experts: int,
+    shared_outer: bool,
+    max_loras: int,
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """
+    Returns virtual topk_ids, token_lora_mask, and virtual_num_experts.
+    """
+    M, top_k = topk_ids.shape
+    device = topk_ids.device
+
+    if shared_outer:
+        num_experts_for_weight = 1
+        # For shared_outer, we need topk_ids to be zeros
+        zero_topk = torch.zeros_like(topk_ids)
+        input_topk = zero_topk
+    else:
+        num_experts_for_weight = num_experts
+        input_topk = topk_ids
+
+    virtual_topk_ids = torch.empty_like(topk_ids)
+    token_lora_mask = torch.empty(M, dtype=torch.bool, device=device)
+
+    BLOCK_SIZE = 1024
+    grid = ((M * top_k + BLOCK_SIZE - 1) // BLOCK_SIZE,)
+
+    _fused_virtual_topk_ids_kernel[grid](
+        input_topk,
+        token_lora_mapping,
+        virtual_topk_ids,
+        token_lora_mask,
+        num_experts_for_weight,
+        M,
+        top_k,
+        BLOCK_SIZE,
+    )
+
+    virtual_num_experts = num_experts_for_weight * max_loras
+    return virtual_topk_ids, token_lora_mask, virtual_num_experts
+
+
+@triton.jit
+def _fused_sanitize_expert_ids_kernel(
+    expert_ids_ptr,
+    output_ptr,
+    num_virtual_experts,
+    N,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    valid = offs < N
+
+    eid = tl.load(expert_ids_ptr + offs, mask=valid, other=0)
+    result = tl.where(eid < num_virtual_experts, eid, -1)
+    tl.store(output_ptr + offs, result, mask=valid)
+
+
+def fused_sanitize_expert_ids(
+    expert_ids: torch.Tensor,
+    num_virtual_experts: int,
+) -> torch.Tensor:
+    """
+    Sanitize expert_ids by replacing values >= num_virtual_experts with -1.
+
+    Returns a new tensor with expert_ids >= num_virtual_experts replaced by -1.
+    """
+    N = expert_ids.numel()
+    output = torch.empty_like(expert_ids)
+
+    BLOCK_SIZE = 1024
+    grid = ((N + BLOCK_SIZE - 1) // BLOCK_SIZE,)
+
+    _fused_sanitize_expert_ids_kernel[grid](
+        expert_ids,
+        output,
+        num_virtual_experts,
+        N,
+        BLOCK_SIZE,
+    )
+    return output
+
+
+@triton.jit
+def _moe_lora_shrink_splitk_kernel(
+    # Pointers
+    a_ptr,  # type: ignore  # [num_tokens, K]
+    b_ptr,  # type: ignore  # [num_virtual_experts, N, K]
+    c_ptr,  # type: ignore  # [num_tokens * top_k, N]  (pre-zeroed when SPLIT_K > 1)
+    sorted_token_ids_ptr,  # type: ignore
+    expert_ids_ptr,  # type: ignore
+    num_tokens_post_padded_ptr,  # type: ignore
+    # Dimensions
+    N,  # type: ignore
+    K,  # type: ignore
+    num_valid_tokens,  # type: ignore
+    # Strides
+    stride_am,  # type: ignore
+    stride_ak,  # type: ignore
+    stride_be,  # type: ignore
+    stride_bn,  # type: ignore
+    stride_bk,  # type: ignore
+    stride_cm,  # type: ignore
+    stride_cn,  # type: ignore
+    # Constexprs
+    top_k: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+):
+    """Split-K grouped GEMM for the LoRA A (shrink) stage with few virtual experts."""
+    pid = tl.program_id(0)
+    pid_sk = pid % SPLIT_K
+    pid_mn = pid // SPLIT_K
+
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    num_pid_m = tl.cdiv(num_tokens_post_padded, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid_mn // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid_mn % num_pid_in_group) % group_size_m)
+    pid_n = (pid_mn % num_pid_in_group) // group_size_m
+
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    # Token routing (same pattern as fused_moe_triton_kernels)
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+
+    off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+    if off_expert == -1:
+        return
+
+    # Pointers
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+
+    a_ptrs = a_ptr + (
+        offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
+    )
+    b_ptrs = (
+        b_ptr
+        + off_expert * stride_be
+        + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    )
+
+    # Accumulate
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
+    for k in range(0, grid_k):
+        k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
+        k_mask = offs_k[:, None] < k_remaining
+        a = tl.load(
+            a_ptrs,
+            mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
+            other=0.0,
+        )
+        b = tl.load(b_ptrs, mask=k_mask, other=0.0)
+        accumulator += tl.dot(a, b.to(a.dtype))
+        a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
+
+    accumulator = accumulator.to(c_ptr.dtype.element_ty)
+
+    # Write output
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+    if SPLIT_K == 1:
+        tl.store(c_ptrs, accumulator, mask=c_mask)
+    else:
+        tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")
+
+
+def _invoke_moe_lora_shrink_splitk(
+    hidden_states: torch.Tensor,
+    weight: torch.Tensor,
+    output: torch.Tensor,
+    topk_ids: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    top_k: int,
+    config: dict[str, Any],
+) -> None:
+    """Launch split-K shrink kernel for LoRA A with few virtual experts."""
+    N = weight.shape[1]
+    K = weight.shape[2]
+    BLOCK_SIZE_M = config["BLOCK_SIZE_M"]
+    BLOCK_SIZE_N = min(config.get("BLOCK_SIZE_N", 64), max(16, N))
+    BLOCK_SIZE_K = config.get("BLOCK_SIZE_K", 64)
+    GROUP_SIZE_M = config.get("GROUP_SIZE_M", 1)
+
+    num_m_blocks = triton.cdiv(sorted_token_ids.shape[0], BLOCK_SIZE_M)
+    num_n_blocks = triton.cdiv(N, BLOCK_SIZE_N)
+    base_grid = num_m_blocks * num_n_blocks
+    max_split_k = max(1, K // BLOCK_SIZE_K)
+    SPLIT_K = min(max_split_k, max(1, 128 // base_grid)) if base_grid < 128 else 1
+
+    grid = (SPLIT_K * base_grid,)
+
+    _moe_lora_shrink_splitk_kernel[grid](
+        hidden_states,
+        weight,
+        output,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        topk_ids.numel(),
+        hidden_states.stride(0),
+        hidden_states.stride(1),
+        weight.stride(0),
+        weight.stride(1),
+        weight.stride(2),
+        output.stride(0),
+        output.stride(1),
+        top_k=top_k,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        GROUP_SIZE_M=GROUP_SIZE_M,
+        SPLIT_K=SPLIT_K,
+        num_warps=config.get("num_warps", 4),
+        num_stages=config.get("num_stages", 4),
+    )
+
+
+def _align_block_size_jit(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """CUDA JIT align_block_size for num_experts > 1024 (up to 8191).
+
+    Uses the v2 kernel from moe_align_kernel.cu which supports large expert
+    counts via per-thread multi-expert processing and a two-level warp scan,
+    replacing the previous pure-PyTorch fallback that had excessive CPU overhead
+    from 15+ individual kernel launches and torch.argsort.
+
+    The JIT kernel uses a +1 offset convention: topk_ids are shifted by +1 so
+    that the EP sentinel value (-1) maps to bucket 0. The kernel internally
+    handles histogram, padded prefix-sum, expert_ids assignment, and token
+    scattering in just 2–3 CUDA kernel launches.
+    """
+    assert num_experts <= 8191, (
+        f"_align_block_size_jit supports at most 8191 experts "
+        f"(num_moe_experts * max_loras), got {num_experts}"
+    )
+
+    device = topk_ids.device
+    flat_topk_ids = topk_ids.reshape(-1)
+    if flat_topk_ids.dtype == torch.int64:
+        flat_topk_ids = flat_topk_ids.to(torch.int32)
+    num_total_tokens = flat_topk_ids.numel()
+
+    if num_total_tokens == 0:
+        empty = torch.empty(0, dtype=torch.int32, device=device)
+        return empty, empty, torch.zeros(1, dtype=torch.int32, device=device)
+
+    # JIT kernel uses +1 offset convention: -1 -> bucket 0 (sentinel),
+    # expert i -> bucket i+1. So pass num_experts + 1 as the bucket count.
+    jit_num_experts = num_experts + 1
+
+    if num_total_tokens < jit_num_experts:
+        max_num_tokens_padded = num_total_tokens * block_size
+    else:
+        max_num_tokens_padded = num_total_tokens + jit_num_experts * (block_size - 1)
+
+    # Align every sub-buffer offset to a multiple of 4 (VEC_SIZE). The CUDA
+    # kernel fills sorted_token_ids with vectorized int4 writes whose last
+    # store can spill up to 3 int32s past the logical end. With a fused
+    # allocation the spill would corrupt the adjacent sub-buffer.
+    _A4 = lambda n: (n + 3) & ~3  # noqa: E731
+    max_num_tokens_padded = _A4(max_num_tokens_padded)
+    max_num_m_blocks = (max_num_tokens_padded + block_size - 1) // block_size
+    max_num_m_blocks_padded = _A4(max_num_m_blocks)
+    num_post_pad_size = _A4(1)  # 1 element, padded to 4
+    cumsum_size = _A4(jit_num_experts + 1)
+
+    # Single allocation sliced into 4 views (zero-copy) to avoid
+    # per-call Python overhead of 4 separate torch.empty calls.
+    total_buf = (
+        max_num_tokens_padded
+        + max_num_m_blocks_padded
+        + num_post_pad_size
+        + cumsum_size
+    )
+    buf = torch.empty(total_buf, dtype=torch.int32, device=device)
+    off = 0
+    sorted_token_ids = buf[off : off + max_num_tokens_padded]
+    off += max_num_tokens_padded
+    expert_ids = buf[off : off + max_num_m_blocks]
+    off += max_num_m_blocks_padded
+    num_tokens_post_padded = buf[off : off + 1]
+    off += num_post_pad_size
+    cumsum_buffer = buf[off : off + jit_num_experts + 1]
+
+    jit_moe_align_block_size(
+        flat_topk_ids,
+        jit_num_experts,
+        block_size,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        cumsum_buffer,
+        True,  # pad_sorted_token_ids
+    )
+
+    return sorted_token_ids, expert_ids, num_tokens_post_padded
+
+
+@torch.compile(dynamic=True)
+def _align_block_size_torch(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Pure-PyTorch align_block_size for num_experts > 1024, compiled via torch.compile.
+
+    Fallback for platforms where the CUDA JIT kernel is unavailable (e.g. AMD/ROCm).
+
+    Out-of-range topk_ids (negative sentinels left by EP dispatch, or virtual-
+    expert IDs >= num_experts produced when those sentinels are combined with
+    a per-adapter offset) are routed into a dedicated sentinel bucket. Without
+    this, indexing ``padded_offsets[sorted_expert_ids]`` would wrap (-1) or
+    OOB-read, and the bad expert ids would propagate into the downstream LoRA
+    GEMM as real expert slots.
+    """
+    device = topk_ids.device
+    flat_topk_ids = topk_ids.reshape(-1).to(torch.int64)
+    num_total_tokens = flat_topk_ids.numel()
+
+    sentinel = num_experts
+    valid_mask = (flat_topk_ids >= 0) & (flat_topk_ids < num_experts)
+    safe_topk_ids = torch.where(
+        valid_mask,
+        flat_topk_ids,
+        torch.full_like(flat_topk_ids, sentinel),
+    )
+
+    bucket_count = num_experts + 1
+    max_total_padded_tokens = (
+        (num_total_tokens + bucket_count * (block_size - 1) + block_size - 1)
+        // block_size
+    ) * block_size
+    max_num_blocks = max_total_padded_tokens // block_size
+
+    sorted_token_ids = torch.full(
+        (max_total_padded_tokens,),
+        num_total_tokens,
+        dtype=torch.int32,
+        device=device,
+    )
+    expert_ids = torch.full(
+        (max_num_blocks,),
+        -1,
+        dtype=torch.int32,
+        device=device,
+    )
+
+    if num_total_tokens == 0:
+        num_tokens_post_padded = torch.zeros((1,), dtype=torch.int32, device=device)
+        return sorted_token_ids, expert_ids, num_tokens_post_padded
+
+    sorted_order = torch.argsort(safe_topk_ids)
+    sorted_expert_ids = safe_topk_ids[sorted_order]
+    expert_range = torch.arange(bucket_count, device=device, dtype=torch.int64)
+    counts_offsets = torch.searchsorted(sorted_expert_ids, expert_range, right=False)
+    counts_end = torch.searchsorted(sorted_expert_ids, expert_range, right=True)
+    counts = counts_end - counts_offsets
+    padded_counts = ((counts + block_size - 1) // block_size) * block_size
+    total_padded_tokens = padded_counts.sum().to(torch.int32).reshape(1)
+    padded_offsets = torch.cumsum(padded_counts, dim=0) - padded_counts
+
+    token_ranks = (
+        torch.arange(num_total_tokens, device=device, dtype=torch.int64)
+        - counts_offsets[sorted_expert_ids]
+    )
+    output_positions = padded_offsets[sorted_expert_ids] + token_ranks
+    sorted_token_ids.scatter_(
+        0,
+        output_positions.to(torch.int64),
+        sorted_order.to(torch.int32),
+    )
+
+    block_counts = padded_counts // block_size
+    real_block_counts = block_counts.clone()
+    real_block_counts[sentinel] = 0
+    actual_num_blocks = real_block_counts.sum()
+
+    if max_num_blocks <= 0:
+        return sorted_token_ids, expert_ids, total_padded_tokens
+
+    block_offsets = torch.cumsum(real_block_counts, dim=0)
+    all_block_positions = torch.arange(max_num_blocks, device=device, dtype=torch.int64)
+    assigned_experts = torch.searchsorted(
+        block_offsets, all_block_positions, right=True
+    ).to(torch.int32)
+    expert_ids.copy_(
+        torch.where(
+            all_block_positions < actual_num_blocks,
+            assigned_experts,
+            torch.full_like(assigned_experts, -1),
+        )
+    )
+
+    return sorted_token_ids, expert_ids, total_padded_tokens
+
+
+def _align_block_size_large(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Dispatch to the CUDA JIT kernel when available, otherwise fall back to
+    the pure-PyTorch torch.compile path (needed on AMD/ROCm or when the JIT
+    module fails to load)."""
+    try:
+        return _align_block_size_jit(topk_ids, block_size, num_experts)
+    except Exception:
+        return _align_block_size_torch(topk_ids, block_size, num_experts)
+
+
+def _merged_experts_fused_moe_lora_add_fake(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+) -> None:
+    return
+
+
+def _merged_experts_fused_moe_lora_add_impl(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+    routing_cache: dict | None = None,
+) -> None:
+    """
+    1. Prepare virtual expert routing metadata from topk_ids + token_lora_mapping * num_experts.
+    2. Flatten LoRA weights from [max_loras, num_experts, ...] to [max_loras * num_experts, ...].
+    3. Run regular SGLang fused-MoE kernels for LoRA A and LoRA B.
+    4. Mask out tokens with token_lora_mapping == -1 on the add path.
+    """
+    max_loras, _, max_lora_rank, _ = lora_a.shape
+    input_top_k = 1 if hidden_states.shape[0] == topk_ids.numel() else topk_ids.shape[1]
+
+    def _merge_lora_expert_weight(t: torch.Tensor) -> torch.Tensor:
+        # [max_loras, num_experts, x, y] -> [max_loras * num_experts, x, y]
+        return t.reshape(t.shape[0] * t.shape[1], t.shape[2], t.shape[3])
+
+    def _get_stage_config(
+        weight: torch.Tensor,
+        stage_top_k: int,
+    ) -> dict[str, Any]:
+        from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config import (
+            get_config_dtype_str,
+            try_get_optimal_moe_config,
+        )
+
+        config_dtype = get_config_dtype_str(dtype=hidden_states.dtype)
+        get_config_func = functools.partial(
+            try_get_optimal_moe_config,
+            weight.shape,
+            weight.shape,
+            stage_top_k,
+            config_dtype,
+        )
+        try:
+            cfg = get_config_func(token_lora_mapping.shape[0])
+        except ValueError:
+            K_dim = weight.shape[2]
+            N_dim = weight.shape[1]
+            if K_dim >= 1024:
+                default_block_k = 256
+            elif K_dim >= 64:
+                default_block_k = 64
+            else:
+                default_block_k = max(16, K_dim)
+            cfg = {
+                "BLOCK_SIZE_M": 64,
+                "BLOCK_SIZE_N": min(64, max(16, N_dim)),
+                "BLOCK_SIZE_K": min(default_block_k, max(16, K_dim)),
+                "GROUP_SIZE_M": 1,
+                "num_warps": 4,
+                "num_stages": 4,
+            }
+        return cfg
+
+    def _align_block_size(
+        topk_ids: torch.Tensor,
+        block_size: int,
+        num_experts: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # The native align kernel consumes num_experts + 1 internally for its
+        # sentinel bucket, so the 1024-expert boundary must use the fallback path.
+        if num_experts < 1024:
+            from sglang.srt.layers.moe.moe_runner.triton_utils.moe_align_block_size import (
+                moe_align_block_size as native_moe_align_block_size,
+            )
+
+            return native_moe_align_block_size(topk_ids, block_size, num_experts)
+        return _align_block_size_large(topk_ids, block_size, num_experts)
+
+    def _get_routing(
+        topk_ids: torch.Tensor,
+        token_lora_mapping: torch.Tensor,
+        num_experts: int,
+        shared_outer: bool,
+        block_size: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Check routing_cache for cross-call reuse (gate_up and down share routing)
+        cache_key = (num_experts, shared_outer, block_size)
+        if routing_cache is not None:
+            cached = routing_cache.get(cache_key)
+            if cached is not None:
+                return cached
+
+        virtual_topk_ids, token_lora_mask, virtual_num_experts = (
+            _fused_virtual_topk_ids(
+                topk_ids, token_lora_mapping, num_experts, shared_outer, max_loras
+            )
+        )
+        sorted_token_ids, expert_ids, num_tokens_post_padded = _align_block_size(
+            virtual_topk_ids,
+            block_size=block_size,
+            num_experts=virtual_num_experts,
+        )
+        # _align_block_size uses a worst-case padded allocation. Trim the routing buffers
+        # to a tighter upper bound so we keep the real routed work but drop unused padding
+        num_tokens = topk_ids.numel()
+        max_nonempty = min(num_tokens, virtual_num_experts)
+        tight_padded = (
+            triton.cdiv(num_tokens + max_nonempty * (block_size - 1), block_size)
+            * block_size
+        )
+        sorted_token_ids = sorted_token_ids[:tight_padded]
+        expert_ids = expert_ids[: tight_padded // block_size]
+        expert_ids = fused_sanitize_expert_ids(expert_ids, virtual_num_experts)
+        result = (
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            token_lora_mask,
+        )
+
+        if routing_cache is not None:
+            routing_cache[cache_key] = result
+
+        return result
+
+    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+        invoke_fused_moe_kernel,
+    )
+
+    lora_a_virtual = _merge_lora_expert_weight(lora_a)
+    lora_b_virtual = _merge_lora_expert_weight(lora_b)
+    num_experts_a = lora_a.shape[1]
+    num_experts_b = lora_b.shape[1]
+
+    intermediate = torch.zeros(
+        [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank],
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+    a_stage_config = _get_stage_config(lora_a_virtual, input_top_k)
+    (
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mask,
+    ) = _get_routing(
+        topk_ids,
+        token_lora_mapping,
+        num_experts_a,
+        experts_shared_outer_loras_a,
+        a_stage_config["BLOCK_SIZE_M"],
+    )
+
+    _invoke_moe_lora_shrink_splitk(
+        hidden_states,
+        lora_a_virtual,
+        intermediate.view(-1, max_lora_rank),
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        input_top_k,
+        a_stage_config,
+    )
+
+    b_stage_config = _get_stage_config(lora_b_virtual, 1)
+    (
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mask,
+    ) = _get_routing(
+        topk_ids,
+        token_lora_mapping,
+        num_experts_b,
+        experts_shared_outer_loras_b,
+        b_stage_config["BLOCK_SIZE_M"],
+    )
+
+    invoke_fused_moe_kernel(
+        intermediate.view(-1, max_lora_rank),
+        lora_b_virtual,
+        None,
+        output,
+        None,
+        None,
+        None,
+        topk_weights,
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        mul_routed_weight,
+        1,
+        b_stage_config,
+        tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16,
+        False,
+        False,
+        False,
+        False,
+        False,
+        None,
+        fuse_add_to_output=True,
+        add_output_mask=token_lora_mask,
+        router_topk=topk_ids.shape[1],
+    )
+
+
+def _merged_experts_fused_moe_lora_add_op(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+) -> None:
+    _merged_experts_fused_moe_lora_add_impl(
+        output,
+        hidden_states,
+        lora_a,
+        lora_b,
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        mul_routed_weight,
+        experts_shared_outer_loras_a,
+        experts_shared_outer_loras_b,
+    )
+
+
+from sglang.srt.utils.common import direct_register_custom_op
+
+direct_register_custom_op(
+    op_name="merged_experts_fused_moe_lora_add",
+    op_func=_merged_experts_fused_moe_lora_add_op,
+    mutates_args=["output"],
+    fake_impl=_merged_experts_fused_moe_lora_add_fake,
+)
+
+
+def merged_experts_fused_moe_lora_add(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+    routing_cache: dict | None = None,
+) -> None:
+    """Public API: wraps the registered op with routing_cache support."""
+    _merged_experts_fused_moe_lora_add_impl(
+        output,
+        hidden_states,
+        lora_a,
+        lora_b,
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        mul_routed_weight,
+        experts_shared_outer_loras_a,
+        experts_shared_outer_loras_b,
+        routing_cache,
+    )

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",
+]

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"]

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"]

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.<module>``); 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__ = []

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

@@ -0,0 +1,418 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import (
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_xpu,
+    next_power_of_2,
+)
+
+_is_cpu = is_cpu()
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_npu = is_npu()
+_is_musa = is_musa()
+_is_xpu = is_xpu()
+
+if _is_cpu:
+    from sgl_kernel import assign_extend_cache_locs_cpu, assign_req_to_token_pool_cpu
+
+
+@triton.jit
+def assign_req_to_token_pool(
+    req_pool_indices,
+    req_to_token,
+    start_offset,
+    end_offset,
+    out_cache_loc,
+    pool_len: tl.constexpr,
+    bs_upper: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 32
+    pid = tl.program_id(axis=0)
+    kv_start = tl.load(start_offset + pid)
+    kv_end = tl.load(end_offset + pid)
+    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
+
+    length_offset = tl.arange(0, bs_upper)
+    start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0)
+    end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0)
+    out_offset = tl.sum(end - start, axis=0)
+
+    out_cache_ptr = out_cache_loc + out_offset
+
+    save_offset = tl.arange(0, BLOCK_SIZE) + kv_start
+    load_offset = tl.arange(0, BLOCK_SIZE)
+
+    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+    for _ in range(num_loop):
+        mask = save_offset < kv_end
+        data = tl.load(out_cache_ptr + load_offset, mask=mask)
+        tl.store(token_pool + save_offset, data, mask=mask)
+        save_offset += BLOCK_SIZE
+        load_offset += BLOCK_SIZE
+
+
+def assign_req_to_token_pool_func(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    batch_size: int,
+):
+    if _is_cpu:
+        assign_req_to_token_pool_cpu(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+        )
+        return
+    assign_req_to_token_pool[(batch_size,)](
+        req_pool_indices,
+        req_to_token,
+        start_offset,
+        end_offset,
+        out_cache_loc,
+        req_to_token.shape[1],
+        next_power_of_2(batch_size),
+    )
+
+
+@triton.jit
+def assign_draft_cache_locs_contiguous(
+    req_pool_indices,
+    req_to_token,
+    seq_lens,
+    out_cache_loc,
+    pool_len: tl.constexpr,
+    topk: tl.constexpr,
+    speculative_num_steps: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 128
+    pid = tl.program_id(axis=0)
+
+    copy_len = topk * speculative_num_steps
+    out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps
+
+    # Copy from req_to_token to out_cache_loc
+    kv_start = tl.load(seq_lens + pid)
+    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
+    num_loop = tl.cdiv(copy_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = copy_offset < copy_len
+        data = tl.load(token_pool + kv_start + copy_offset, mask=mask)
+        tl.store(out_cache_ptr + copy_offset, data, mask=mask)
+
+
+@triton.jit
+def generate_draft_decode_kv_indices(
+    req_pool_indices,
+    req_to_token,
+    paged_kernel_lens,
+    kv_indices,
+    kv_indptr,
+    positions,
+    pool_len: tl.constexpr,
+    kv_indices_stride: tl.constexpr,
+    kv_indptr_stride: tl.constexpr,
+    bs_upper: tl.constexpr,
+    iter_upper: tl.constexpr,
+    num_tokens_upper: tl.constexpr,
+    page_size: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 128
+    iters = tl.program_id(axis=0)
+    bid = tl.program_id(axis=1)
+    topk_id = tl.program_id(axis=2)
+
+    num_steps = tl.num_programs(axis=0)
+    num_seqs = tl.num_programs(axis=1)
+    topk = tl.num_programs(axis=2)
+
+    kv_indices += kv_indices_stride * iters
+    kv_indptr += kv_indptr_stride * iters
+    iters += 1
+
+    load_offset = tl.arange(0, bs_upper)
+    seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0)
+    seq_len = tl.load(paged_kernel_lens + bid)
+    cum_seq_len = tl.sum(seq_lens)
+
+    # Update kv_indices
+    kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters)
+    kv_ptr = kv_indices + kv_offset
+    token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len
+
+    kv_offset = tl.arange(0, BLOCK_SIZE)
+    num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
+    for _ in range(num_loop):
+        mask = kv_offset < seq_len
+        data = tl.load(token_pool_ptr + kv_offset, mask=mask)
+        tl.store(kv_ptr + kv_offset, data, mask=mask)
+        kv_offset += BLOCK_SIZE
+
+    extend_offset = tl.arange(0, iter_upper)
+    if page_size == 1 or topk == 1:
+        extend_data = tl.load(
+            token_pool_ptr + seq_len + topk_id * num_steps + tl.arange(0, iter_upper),
+            mask=extend_offset < iters,
+        )
+    else:
+        prefix_len = seq_len
+        last_page_len = prefix_len % page_size
+        num_new_pages_per_topk = (
+            last_page_len + num_steps + page_size - 1
+        ) // page_size
+        prefix_base = seq_len // page_size * page_size
+        start = (
+            prefix_base + topk_id * num_new_pages_per_topk * page_size + last_page_len
+        )
+        extend_data = tl.load(
+            token_pool_ptr + start + extend_offset,
+            mask=extend_offset < iters,
+        )
+
+    tl.store(kv_ptr + seq_len + extend_offset, extend_data, mask=extend_offset < iters)
+
+    # Update kv_indptr
+    bs_offset = tl.arange(0, num_tokens_upper)
+
+    zid = bid * topk + topk_id
+    if zid == 0:
+        zid = num_seqs * topk
+    positions = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0)
+    base = tl.sum(positions)
+    tl.store(kv_indptr + zid, base + zid * iters)
+
+
+@triton.jit
+def align_evict_mask_to_page_size(
+    seq_lens,
+    evict_mask,
+    page_size: tl.constexpr,
+    num_draft_tokens: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    t_range = tl.arange(0, BLOCK_SIZE)
+
+    bid = tl.program_id(axis=0)
+    seq_len = tl.load(seq_lens + bid)
+    io_mask = t_range < num_draft_tokens
+    mask_row = tl.load(
+        evict_mask + bid * num_draft_tokens + t_range, mask=io_mask, other=0
+    )
+
+    num_trues = tl.sum(mask_row)
+    num_false = num_draft_tokens - num_trues
+
+    start = (seq_len + num_false - 1) // page_size * page_size - seq_len
+    for i in range(max(start, 0), min(start + page_size, num_draft_tokens)):
+        tl.store(evict_mask + bid * num_draft_tokens + i, False)
+
+
+@torch.compile(dynamic=True, disable=_is_npu)
+def get_src_tgt_cache_loc(
+    seq_lens: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    accept_index: torch.Tensor,
+    num_correct_drafts: torch.Tensor,
+    draft_token_num: int,
+    page_size: int,
+):
+    src_cache_loc = out_cache_loc[accept_index]
+    # zeros_like, not empty_like: any uncovered tail stays at slot 0 (padding)
+    # instead of caching-allocator garbage.
+    tgt_cache_loc = torch.zeros_like(src_cache_loc)
+    extended_len = seq_lens + draft_token_num
+    keep_len = torch.minimum(
+        (seq_lens + num_correct_drafts + 1 + page_size - 1) // page_size * page_size,
+        extended_len,
+    )
+    to_free_num_slots = extended_len - keep_len
+    return src_cache_loc, tgt_cache_loc, to_free_num_slots
+
+
+@triton.jit
+def get_target_cache_loc(
+    tgt_cache_loc,
+    to_free_slots,
+    num_correct_drafts,
+    to_free_num_slots,
+    out_cache_loc,
+    num_verify_tokens: tl.constexpr,
+    num_verify_tokens_upper: tl.constexpr,
+    bs_upper: tl.constexpr,
+):
+    bid = tl.program_id(axis=0)
+    offset = tl.arange(0, num_verify_tokens_upper)
+    bs_offset = tl.arange(0, bs_upper)
+
+    # write the first part to tgt_cache_loc
+    accept_len_all = tl.load(num_correct_drafts + bs_offset, mask=bs_offset < bid)
+    tgt_cache_loc_start = tl.sum(accept_len_all) + bid
+    copy_len = tl.load(num_correct_drafts + bid) + 1
+    out_cache_loc_row = tl.load(
+        out_cache_loc + bid * num_verify_tokens + offset, mask=offset < copy_len
+    )
+    tl.store(
+        tgt_cache_loc + tgt_cache_loc_start + offset,
+        out_cache_loc_row,
+        mask=offset < copy_len,
+    )
+
+    # write the second part to to_free_num_pages
+    to_free_num_slots_all = tl.load(to_free_num_slots + bs_offset, mask=bs_offset < bid)
+    to_free_num_slots_cur = tl.load(to_free_num_slots + bid)
+    out_cache_loc_start = num_verify_tokens - to_free_num_slots_cur
+    to_free_slots_start = tl.sum(to_free_num_slots_all)
+
+    copy_len = to_free_num_slots_cur
+    out_cache_loc_row = tl.load(
+        out_cache_loc + bid * num_verify_tokens + out_cache_loc_start + offset,
+        mask=offset < copy_len,
+    )
+    tl.store(
+        to_free_slots + to_free_slots_start + offset,
+        out_cache_loc_row,
+        mask=offset < copy_len,
+    )
+
+
+@triton.jit
+def filter_finished_cache_loc_kernel(
+    out_cache_loc,
+    tgt_cache_loc,
+    num_correct_drafts,
+    num_accept_tokens_filter,
+    bs_upper: tl.constexpr,
+    num_verify_tokens_upper: tl.constexpr,
+):
+    bid = tl.program_id(0)
+    bs_offset = tl.arange(0, bs_upper)
+
+    num_correct_drafts_all = tl.load(
+        num_correct_drafts + bs_offset, mask=bs_offset < bid
+    )
+    old_start = tl.sum(num_correct_drafts_all) + bid
+
+    num_accept_tokens_filter_all = tl.load(
+        num_accept_tokens_filter + bs_offset, mask=bs_offset < bid
+    )
+    new_start = tl.sum(num_accept_tokens_filter_all)
+
+    copy_len = tl.load(num_accept_tokens_filter + bid)
+    copy_offset = tl.arange(0, num_verify_tokens_upper)
+    value = tl.load(
+        tgt_cache_loc + old_start + copy_offset, mask=copy_offset < copy_len
+    )
+    tl.store(
+        out_cache_loc + new_start + copy_offset, value, mask=copy_offset < copy_len
+    )
+
+
+@triton.jit
+def assign_extend_cache_locs(
+    req_pool_indices,
+    req_to_token,
+    start_offset,
+    end_offset,
+    out_cache_loc,
+    pool_len: tl.constexpr,
+    bs_upper: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 32
+    pid = tl.program_id(axis=0)
+    kv_start = tl.load(start_offset + pid)
+    kv_end = tl.load(end_offset + pid)
+    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
+
+    length_offset = tl.arange(0, bs_upper)
+    start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0)
+    end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0)
+    out_offset = tl.sum(end - start, axis=0)
+
+    out_cache_ptr = out_cache_loc + out_offset
+
+    load_offset = tl.arange(0, BLOCK_SIZE) + kv_start
+    save_offset = tl.arange(0, BLOCK_SIZE)
+
+    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+    for _ in range(num_loop):
+        mask = load_offset < kv_end
+        data = tl.load(token_pool + load_offset, mask=mask)
+        tl.store(out_cache_ptr + save_offset, data, mask=mask)
+        load_offset += BLOCK_SIZE
+        save_offset += BLOCK_SIZE
+
+
+def assign_extend_cache_locs_func(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    batch_size: int,
+    draft_token_num: int,
+    device,
+) -> torch.Tensor:
+    if _is_cuda or _is_hip or _is_musa or _is_xpu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int64,
+            device=device,
+        )
+        assign_extend_cache_locs[(batch_size,)](
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+            next_power_of_2(batch_size),
+        )
+
+        return out_cache_loc
+
+    elif _is_npu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int32,
+            device=device,
+        )
+        torch.ops.npu.cache_loc_update(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+        )
+
+        return out_cache_loc
+
+    elif _is_cpu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int64,
+            device=device,
+        )
+        assign_extend_cache_locs_cpu(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+        )
+
+        return out_cache_loc

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

@@ -0,0 +1,246 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _dflash_accept_bonus_contig_kernel(
+    candidates_ptr,
+    target_top1_ptr,
+    accept_lens_out_ptr,
+    commit_lens_out_ptr,
+    bonus_ids_out_ptr,
+    out_tokens_ptr,
+    prefix_lens_ptr,
+    new_seq_lens_out_ptr,
+    candidates_row_stride,
+    target_row_stride,
+    accept_stride,
+    commit_stride,
+    bonus_stride,
+    out_tokens_row_stride,
+    prefix_lens_stride,
+    new_seq_lens_stride,
+    block_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_SIZE)
+    row_mask = cols < block_size
+    draft_mask = cols < (block_size - 1)
+
+    candidate_row_ptr = candidates_ptr + row * candidates_row_stride
+    target_row_ptr = target_top1_ptr + row * target_row_stride
+    candidate_tail = tl.load(candidate_row_ptr + cols + 1, mask=draft_mask, other=0)
+
+    accept_len = tl.full((), 0, tl.int32)
+    prefix_live = tl.full((), 1, tl.int32)
+    for col in range(BLOCK_SIZE - 1):
+        in_range = col < (block_size - 1)
+        candidate_id = tl.load(candidate_row_ptr + (col + 1), mask=in_range, other=0)
+        target_id = tl.load(target_row_ptr + col, mask=in_range, other=0)
+        match_i32 = (candidate_id == target_id).to(tl.int32)
+        keep = in_range & (prefix_live != 0) & (match_i32 != 0)
+        accept_len += keep.to(tl.int32)
+        prefix_live = tl.where(in_range, prefix_live & match_i32, prefix_live)
+
+    commit_len = accept_len + 1
+    bonus_id = tl.load(target_row_ptr + accept_len.to(tl.int64))
+    new_seq_len = tl.load(prefix_lens_ptr + row * prefix_lens_stride) + commit_len
+
+    tl.store(accept_lens_out_ptr + row * accept_stride, accept_len)
+    tl.store(commit_lens_out_ptr + row * commit_stride, commit_len)
+    tl.store(bonus_ids_out_ptr + row * bonus_stride, bonus_id)
+    tl.store(new_seq_lens_out_ptr + row * new_seq_lens_stride, new_seq_len)
+
+    out_val = tl.where(draft_mask, candidate_tail, 0)
+    out_val = tl.where(cols == accept_len, bonus_id, out_val)
+    tl.store(
+        out_tokens_ptr + row * out_tokens_row_stride + cols, out_val, mask=row_mask
+    )
+
+
+def _pick_num_warps(block_size: int) -> int:
+    if block_size <= 16:
+        return 1
+    if block_size <= 32:
+        return 2
+    if block_size <= 64:
+        return 4
+    return 8
+
+
+def _is_row_major_contiguous_2d(x: torch.Tensor) -> bool:
+    return x.ndim == 2 and x.is_contiguous()
+
+
+def _compute_dflash_accept_bonus_triton_unchecked(
+    candidates: torch.Tensor,
+    target_top1: torch.Tensor,
+    accept_lens_out: torch.Tensor,
+    commit_lens_out: torch.Tensor,
+    bonus_ids_out: torch.Tensor,
+    out_tokens_out: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    new_seq_lens_out: torch.Tensor,
+) -> None:
+    batch_size, block_size = candidates.shape
+    if batch_size == 0:
+        return
+
+    if not _is_row_major_contiguous_2d(candidates):
+        raise ValueError("DFLASH Triton accept_bonus requires contiguous candidates.")
+    if not _is_row_major_contiguous_2d(target_top1):
+        raise ValueError("DFLASH Triton accept_bonus requires contiguous target_top1.")
+    if not _is_row_major_contiguous_2d(out_tokens_out):
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous out_tokens_out."
+        )
+    if not accept_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous accept_lens_out."
+        )
+    if not commit_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous commit_lens_out."
+        )
+    if not bonus_ids_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous bonus_ids_out."
+        )
+    if prefix_lens.ndim != 1:
+        raise ValueError("DFLASH Triton accept_bonus requires 1D prefix_lens.")
+    if not new_seq_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous new_seq_lens_out."
+        )
+
+    block = triton.next_power_of_2(block_size)
+    num_warps = _pick_num_warps(block)
+    _dflash_accept_bonus_contig_kernel[(batch_size,)](
+        candidates,
+        target_top1,
+        accept_lens_out,
+        commit_lens_out,
+        bonus_ids_out,
+        out_tokens_out,
+        prefix_lens,
+        new_seq_lens_out,
+        candidates.stride(0),
+        target_top1.stride(0),
+        accept_lens_out.stride(0),
+        commit_lens_out.stride(0),
+        bonus_ids_out.stride(0),
+        out_tokens_out.stride(0),
+        prefix_lens.stride(0),
+        new_seq_lens_out.stride(0),
+        block_size,
+        BLOCK_SIZE=block,
+        num_warps=num_warps,
+    )
+
+
+@triton.jit
+def _prepare_dflash_draft_block_contig_kernel(
+    bonus_tokens_ptr,
+    prefix_lens_ptr,
+    req_pool_indices_ptr,
+    req_to_token_ptr,
+    block_ids_out_ptr,
+    positions_out_ptr,
+    cache_loc_out_ptr,
+    bonus_tokens_stride,
+    prefix_lens_stride,
+    req_pool_indices_stride,
+    req_to_token_row_stride,
+    block_ids_row_stride,
+    positions_row_stride,
+    cache_loc_row_stride,
+    req_to_token_width,
+    block_size,
+    mask_token_id,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_SIZE)
+    row_mask = cols < block_size
+
+    prefix_len = tl.load(prefix_lens_ptr + row * prefix_lens_stride)
+    req_idx = tl.load(req_pool_indices_ptr + row * req_pool_indices_stride)
+    bonus_token = tl.load(bonus_tokens_ptr + row * bonus_tokens_stride)
+
+    logical_pos = prefix_len.to(tl.int64) + cols
+    valid = row_mask & (logical_pos < req_to_token_width)
+    req_row_ptr = req_to_token_ptr + req_idx * req_to_token_row_stride
+    slot_ids = tl.load(req_row_ptr + logical_pos, mask=valid, other=0)
+
+    block_ids = tl.full((BLOCK_SIZE,), mask_token_id, tl.int64)
+    block_ids = tl.where(cols == 0, bonus_token.to(tl.int64), block_ids)
+    tl.store(
+        block_ids_out_ptr + row * block_ids_row_stride + cols, block_ids, mask=row_mask
+    )
+    tl.store(
+        positions_out_ptr + row * positions_row_stride + cols,
+        logical_pos,
+        mask=row_mask,
+    )
+    tl.store(
+        cache_loc_out_ptr + row * cache_loc_row_stride + cols,
+        slot_ids.to(tl.int64),
+        mask=row_mask,
+    )
+
+
+def _prepare_dflash_draft_block_unchecked(
+    bonus_tokens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    block_ids_out: torch.Tensor,
+    positions_out: torch.Tensor,
+    cache_loc_out: torch.Tensor,
+    mask_token_id: int,
+) -> None:
+    batch_size = int(bonus_tokens.numel())
+    if batch_size == 0:
+        return
+
+    if req_to_token.ndim != 2 or req_to_token.stride(1) != 1:
+        raise ValueError("DFLASH Triton prepare_block requires row-major req_to_token.")
+    if not _is_row_major_contiguous_2d(block_ids_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous block_ids_out."
+        )
+    if not _is_row_major_contiguous_2d(positions_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous positions_out."
+        )
+    if not _is_row_major_contiguous_2d(cache_loc_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous cache_loc_out."
+        )
+
+    block_size = int(block_ids_out.shape[1])
+    block = triton.next_power_of_2(block_size)
+    num_warps = _pick_num_warps(block)
+    _prepare_dflash_draft_block_contig_kernel[(batch_size,)](
+        bonus_tokens,
+        prefix_lens,
+        req_pool_indices,
+        req_to_token,
+        block_ids_out,
+        positions_out,
+        cache_loc_out,
+        bonus_tokens.stride(0),
+        prefix_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        block_ids_out.stride(0),
+        positions_out.stride(0),
+        cache_loc_out.stride(0),
+        int(req_to_token.shape[1]),
+        block_size,
+        int(mask_token_id),
+        BLOCK_SIZE=block,
+        num_warps=num_warps,
+    )

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

@@ -0,0 +1,91 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cpu, next_power_of_2
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import fill_accept_out_cache_loc_cpu, fill_bonus_tokens_cpu
+
+
+@triton.jit
+def fill_bonus_tokens(
+    accept_tokens,
+    accept_lens,
+    bonus_tokens_ptr,
+    accept_stride: tl.constexpr,
+):
+    # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel
+    # because this kernel reads accept_lens
+    pid = tl.program_id(axis=0)
+    # `accept_lens` includes the bonus token; the last accepted slot is at -1.
+    accept_len = tl.load(accept_lens + pid)
+
+    # accept_stride = per-req width of accept_tokens (= accept_index.shape[1]).
+    bonus_token_idx = accept_stride * pid + accept_len - 1
+    bonus_token = tl.load(accept_tokens + bonus_token_idx)
+    tl.store(bonus_tokens_ptr + pid, bonus_token)
+
+
+def fill_bonus_tokens_func(
+    accept_tokens: torch.Tensor,
+    accept_lens: torch.Tensor,
+    bonus_tokens: torch.Tensor,  # mutable
+    accept_stride: int,
+    batch_size: int,
+):
+    if _is_cpu:
+        fill_bonus_tokens_cpu(
+            accept_tokens,
+            accept_lens,
+            bonus_tokens,
+            accept_stride,
+        )
+        return
+    fill_bonus_tokens[(batch_size,)](
+        accept_tokens,
+        accept_lens,
+        bonus_tokens,
+        accept_stride,
+    )
+
+
+@triton.jit
+def fill_accept_out_cache_loc(
+    accept_index,
+    out_cache_loc,
+    accept_out_cache_loc,
+    size_upper: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    offset = tl.arange(0, size_upper)
+
+    masks = (tl.load(accept_index + offset, offset < pid, other=-1) != -1).to(tl.int64)
+    dst = tl.sum(masks)
+    src = tl.load(accept_index + pid)
+    if src > -1:
+        value = tl.load(out_cache_loc + src)
+        tl.store(accept_out_cache_loc + dst, value)
+
+
+def fill_accept_out_cache_loc_func(
+    accept_index: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    accept_out_cache_loc: torch.Tensor,  # mutable
+    size: int,
+):
+    if _is_cpu:
+        fill_accept_out_cache_loc_cpu(
+            accept_index,
+            out_cache_loc,
+            accept_out_cache_loc,
+        )
+        return
+    fill_accept_out_cache_loc[(size,)](
+        accept_index,
+        out_cache_loc,
+        accept_out_cache_loc,
+        next_power_of_2(size),
+    )

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

@@ -0,0 +1,457 @@
+# 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.
+# ==============================================================================
+"""Fused Triton kernel for DFlash KV materialization.
+
+Combines: KV projection + RMSNorm + RoPE, then pool-managed KV writes.
+"""
+
+from typing import Callable, List, Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fused_norm_rope_kernel_stacked(
+    kv_ptr,  # [total_ctx, n_layers, kv_size * 2]
+    k_norm_weight_ptr,  # [n_layers, head_dim]
+    eps_ptr,  # [n_layers]
+    cos_sin_cache_ptr,  # [max_pos, rotary_dim]
+    positions_ptr,  # [total_ctx]
+    k_out_ptr,  # [n_layers, total_ctx, num_kv_heads, head_dim]
+    v_out_ptr,  # [n_layers, total_ctx, num_kv_heads, head_dim]
+    kv_stride_ctx,
+    kv_stride_layer,
+    k_norm_weight_stride_layer,
+    cos_sin_stride_pos,
+    k_out_stride_layer,
+    k_out_stride_ctx,
+    k_out_stride_head,
+    v_out_stride_layer,
+    v_out_stride_ctx,
+    v_out_stride_head,
+    total_ctx,
+    n_layers: tl.constexpr,
+    num_kv_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    kv_size: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    half_rotary_dim: tl.constexpr,
+    BLOCK_HD: tl.constexpr,
+):
+    """Fused RMSNorm(K) + RoPE(K) materialization. Grid: (total_ctx, num_kv_heads, n_layers)."""
+    ctx_id = tl.program_id(0)
+    head_id = tl.program_id(1)
+    layer_id = tl.program_id(2)
+    if ctx_id >= total_ctx or layer_id >= n_layers:
+        return
+
+    position = tl.load(positions_ptr + ctx_id)
+    eps = tl.load(eps_ptr + layer_id).to(tl.float32)
+    kv_base = kv_ptr + ctx_id * kv_stride_ctx + layer_id * kv_stride_layer
+    k_base = kv_base + head_id * head_dim
+    v_base = kv_base + kv_size + head_id * head_dim
+    k_write = (
+        k_out_ptr
+        + layer_id * k_out_stride_layer
+        + ctx_id * k_out_stride_ctx
+        + head_id * k_out_stride_head
+    )
+    v_write = (
+        v_out_ptr
+        + layer_id * v_out_stride_layer
+        + ctx_id * v_out_stride_ctx
+        + head_id * v_out_stride_head
+    )
+
+    offs = tl.arange(0, BLOCK_HD)
+    mask_hd = offs < head_dim
+    mask_half = offs < half_rotary_dim
+
+    k_raw = tl.load(k_base + offs, mask=mask_hd, other=0.0).to(tl.float32)
+    v_raw = tl.load(v_base + offs, mask=mask_hd, other=0.0)
+
+    inv_rms = tl.rsqrt(tl.sum(k_raw * k_raw) / head_dim + eps)
+    norm_w = tl.load(
+        k_norm_weight_ptr + layer_id * k_norm_weight_stride_layer + offs,
+        mask=mask_hd,
+        other=1.0,
+    ).to(tl.float32)
+    k_normed = k_raw * inv_rms * norm_w
+
+    cos_sin_base = cos_sin_cache_ptr + position * cos_sin_stride_pos
+    cos_v = tl.load(cos_sin_base + offs, mask=mask_half, other=1.0).to(tl.float32)
+    sin_v = tl.load(
+        cos_sin_base + half_rotary_dim + offs, mask=mask_half, other=0.0
+    ).to(tl.float32)
+
+    k_first = tl.where(mask_half, k_normed, 0.0)
+    k_second_raw = tl.load(
+        k_base + half_rotary_dim + offs, mask=mask_half, other=0.0
+    ).to(tl.float32)
+    norm_w_second = tl.load(
+        k_norm_weight_ptr
+        + layer_id * k_norm_weight_stride_layer
+        + half_rotary_dim
+        + offs,
+        mask=mask_half,
+        other=1.0,
+    ).to(tl.float32)
+    k_second = k_second_raw * inv_rms * norm_w_second
+
+    k_rot_first = k_first * cos_v - k_second * sin_v
+    k_rot_second = k_second * cos_v + k_first * sin_v
+
+    tl.store(v_write + offs, v_raw, mask=mask_hd)
+    tl.store(k_write + offs, k_rot_first.to(v_raw.dtype), mask=mask_half)
+    tl.store(
+        k_write + half_rotary_dim + offs, k_rot_second.to(v_raw.dtype), mask=mask_half
+    )
+    mask_pass = (offs >= rotary_dim) & (offs < head_dim)
+    tl.store(k_write + offs, k_normed.to(v_raw.dtype), mask=mask_pass)
+
+
+def _fused_norm_rope_stacked(
+    kv: torch.Tensor,  # [total_ctx, n_layers, kv_size*2]
+    k_norm_weight: torch.Tensor,  # [n_layers, head_dim]
+    eps: torch.Tensor,  # [n_layers]
+    cos_sin_cache: torch.Tensor,  # [max_pos, rotary_dim]
+    positions: torch.Tensor,  # [total_ctx]
+    num_kv_heads: int,
+    head_dim: int,
+    rotary_dim: int,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused RMSNorm + RoPE materialization for all layers."""
+    if kv.ndim != 3:
+        raise ValueError(
+            "Invalid stacked fused KV projection shape: "
+            f"got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2]."
+        )
+
+    total_ctx, n_layers, kv_dim = kv.shape
+    if total_ctx == 0:
+        empty = torch.empty(
+            (n_layers, 0, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device
+        )
+        return empty, empty
+
+    kv_size = num_kv_heads * head_dim
+    if kv_dim != kv_size * 2:
+        raise ValueError(
+            "Invalid fused KV projection shape: "
+            f"got {tuple(kv.shape)}, expected trailing dim {kv_size * 2}."
+        )
+    if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0:
+        raise ValueError(
+            "Invalid fused KV rotary/head dim pair: "
+            f"rotary_dim={rotary_dim}, head_dim={head_dim}."
+        )
+    if k_norm_weight.shape != (n_layers, head_dim):
+        raise ValueError(
+            "Invalid stacked k_norm_weight shape for fused KV materialization: "
+            f"got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}."
+        )
+    if eps.shape != (n_layers,):
+        raise ValueError(
+            "Invalid stacked eps shape for fused KV materialization: "
+            f"got {tuple(eps.shape)}, expected {(n_layers,)}."
+        )
+
+    half_rotary_dim = rotary_dim // 2
+    BLOCK_HD = triton.next_power_of_2(head_dim)
+
+    if positions.device != kv.device:
+        positions = positions.to(device=kv.device, dtype=torch.int64)
+    elif positions.dtype != torch.int64:
+        positions = positions.to(torch.int64)
+
+    expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim)
+    if k_out is None:
+        k_out = torch.empty(expected_shape, dtype=kv.dtype, device=kv.device)
+    else:
+        if k_out.shape != expected_shape:
+            raise ValueError(
+                "Invalid k_out shape for fused KV materialization: "
+                f"got {tuple(k_out.shape)}, expected {expected_shape}."
+            )
+        if k_out.device != kv.device or k_out.dtype != kv.dtype:
+            raise ValueError(
+                "Invalid k_out device/dtype for fused KV materialization: "
+                f"got device={k_out.device}, dtype={k_out.dtype}, "
+                f"expected device={kv.device}, dtype={kv.dtype}."
+            )
+    if v_out is None:
+        v_out = torch.empty_like(k_out)
+    else:
+        if v_out.shape != expected_shape:
+            raise ValueError(
+                "Invalid v_out shape for fused KV materialization: "
+                f"got {tuple(v_out.shape)}, expected {expected_shape}."
+            )
+        if v_out.device != kv.device or v_out.dtype != kv.dtype:
+            raise ValueError(
+                "Invalid v_out device/dtype for fused KV materialization: "
+                f"got device={v_out.device}, dtype={v_out.dtype}, "
+                f"expected device={kv.device}, dtype={kv.dtype}."
+            )
+
+    _fused_norm_rope_kernel_stacked[(total_ctx, num_kv_heads, n_layers)](
+        kv,
+        k_norm_weight,
+        eps,
+        cos_sin_cache,
+        positions,
+        k_out,
+        v_out,
+        kv.stride(0),
+        kv.stride(1),
+        k_norm_weight.stride(0),
+        cos_sin_cache.stride(0),
+        k_out.stride(0),
+        k_out.stride(1),
+        k_out.stride(2),
+        v_out.stride(0),
+        v_out.stride(1),
+        v_out.stride(2),
+        total_ctx,
+        n_layers,
+        num_kv_heads,
+        head_dim,
+        kv_size,
+        rotary_dim,
+        half_rotary_dim,
+        BLOCK_HD,
+    )
+    return k_out, v_out
+
+
+class FusedKVMaterializeHelper:
+    """Fused KV materialization helper using batched projection.
+
+    Uses a single large GEMM across all layers, then a Triton kernel for fused
+    RMSNorm + RoPE materialization across all layers.
+    """
+
+    def __init__(
+        self,
+        layers: List,
+        rotary_emb,
+        num_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        max_position_hint: Optional[int] = None,
+    ):
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.rotary_emb = rotary_emb
+        self.n_layers = len(layers)
+        self.device = device
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.layer_out_dim = 2 * self.kv_size
+
+        self.rotary_dim = int(getattr(rotary_emb, "rotary_dim", head_dim))
+        self.is_neox_style = bool(getattr(rotary_emb, "is_neox_style", True))
+
+        if not self.is_neox_style:
+            raise NotImplementedError("Only neox-style RoPE is supported.")
+        if self.rotary_dim <= 0 or self.rotary_dim > self.head_dim:
+            raise ValueError(
+                "Invalid fused KV rotary/head dim pair: "
+                f"rotary_dim={self.rotary_dim}, head_dim={self.head_dim}."
+            )
+
+        self.max_position_hint = (
+            max(int(max_position_hint) - 1, 0)
+            if max_position_hint is not None
+            else None
+        )
+        self._reserved_rope_cache_len = int(
+            getattr(self.rotary_emb, "cos_sin_cache", torch.empty((0,))).shape[0]
+        )
+        self._mm_out_supported = True
+        self._workspace_capacity = 0
+        self._workspace_dtype: Optional[torch.dtype] = None
+        self._proj_workspace: Optional[torch.Tensor] = None
+        self._k_workspace: Optional[torch.Tensor] = None
+        self._v_workspace: Optional[torch.Tensor] = None
+
+        kv_weights = []
+        k_norm_weights = []
+        eps_values = []
+
+        for layer_id, layer in enumerate(layers):
+            attn = layer.self_attn
+            if int(attn.num_kv_heads) != self.num_kv_heads:
+                raise ValueError(
+                    "num_kv_heads mismatch across layers for fused KV path: "
+                    f"expected {self.num_kv_heads}, got {int(attn.num_kv_heads)} at layer {layer_id}."
+                )
+            if int(attn.head_dim) != self.head_dim:
+                raise ValueError(
+                    "head_dim mismatch across layers for fused KV path: "
+                    f"expected {self.head_dim}, got {int(attn.head_dim)} at layer {layer_id}."
+                )
+            layer_rotary_dim = int(
+                getattr(attn.rotary_emb, "rotary_dim", self.head_dim)
+            )
+            layer_is_neox = bool(getattr(attn.rotary_emb, "is_neox_style", True))
+            if (
+                layer_rotary_dim != self.rotary_dim
+                or layer_is_neox != self.is_neox_style
+            ):
+                raise ValueError(
+                    "RoPE config mismatch across layers for fused KV path: "
+                    f"expected (rotary_dim={self.rotary_dim}, neox={self.is_neox_style}), "
+                    f"got (rotary_dim={layer_rotary_dim}, neox={layer_is_neox}) at layer {layer_id}."
+                )
+
+            qkv_w = attn.qkv_proj.weight
+            kv_weight = qkv_w[attn.q_size : attn.q_size + 2 * attn.kv_size]
+            kv_weights.append(kv_weight)
+            k_norm_weights.append(attn.k_norm.weight)
+            eps_values.append(float(attn.k_norm.variance_epsilon))
+
+        flat_kv_weight = torch.stack(kv_weights).reshape(
+            self.n_layers * self.layer_out_dim, -1
+        )
+        self.flat_kv_weight_t = flat_kv_weight.transpose(0, 1).contiguous()
+        self.k_norm_weights = torch.stack(k_norm_weights).contiguous()
+        self.eps_values = torch.tensor(
+            eps_values, dtype=torch.float32, device=self.device
+        )
+
+        if self.max_position_hint is not None:
+            self._ensure_rope_cache(self.max_position_hint)
+
+    def _ensure_rope_cache(self, max_position: int) -> torch.Tensor:
+        if max_position + 1 > self._reserved_rope_cache_len:
+            ensure_cos_sin_cache_length = getattr(
+                self.rotary_emb, "_ensure_cos_sin_cache_length", None
+            )
+            if callable(ensure_cos_sin_cache_length):
+                ensure_cos_sin_cache_length(max_position)
+                self._reserved_rope_cache_len = int(
+                    self.rotary_emb.cos_sin_cache.shape[0]
+                )
+
+        cos_sin_cache = self.rotary_emb.cos_sin_cache
+        if max_position >= int(cos_sin_cache.shape[0]):
+            raise RuntimeError(
+                "RoPE cos/sin cache is too short for fused KV materialization: "
+                f"max_position={max_position}, cache_len={int(cos_sin_cache.shape[0])}."
+            )
+        if cos_sin_cache.device != self.device:
+            cos_sin_cache = cos_sin_cache.to(self.device)
+        return cos_sin_cache
+
+    def _ensure_workspace(self, total_ctx: int, dtype: torch.dtype) -> None:
+        if (
+            self._workspace_capacity >= total_ctx
+            and self._workspace_dtype == dtype
+            and self._proj_workspace is not None
+            and self._k_workspace is not None
+            and self._v_workspace is not None
+        ):
+            return
+
+        new_capacity = max(1, total_ctx)
+        if self._workspace_capacity > 0:
+            new_capacity = max(new_capacity, self._workspace_capacity * 2)
+
+        self._proj_workspace = torch.empty(
+            (new_capacity, self.n_layers * self.layer_out_dim),
+            dtype=dtype,
+            device=self.device,
+        )
+        self._k_workspace = torch.empty(
+            (self.n_layers, new_capacity, self.num_kv_heads, self.head_dim),
+            dtype=dtype,
+            device=self.device,
+        )
+        self._v_workspace = torch.empty_like(self._k_workspace)
+        self._workspace_capacity = new_capacity
+        self._workspace_dtype = dtype
+
+    def materialize(
+        self,
+        ctx_hidden: torch.Tensor,
+        positions: torch.Tensor,
+        write_layer_kv: Callable[[int, torch.Tensor, torch.Tensor], None],
+    ) -> None:
+        """Materialize KV cache for all layers using batched projection."""
+        total_ctx = ctx_hidden.shape[0]
+        if total_ctx == 0:
+            return
+
+        if positions.ndim != 1:
+            positions = positions.reshape(-1)
+        if positions.numel() != total_ctx:
+            raise ValueError(
+                "positions must match ctx_hidden token count for fused KV materialization: "
+                f"positions={positions.numel()}, total_ctx={total_ctx}."
+            )
+
+        if ctx_hidden.device != self.device:
+            ctx_hidden = ctx_hidden.to(self.device, non_blocking=True)
+        if ctx_hidden.dtype != self.flat_kv_weight_t.dtype:
+            ctx_hidden = ctx_hidden.to(self.flat_kv_weight_t.dtype)
+        if positions.device != self.device:
+            positions = positions.to(
+                device=self.device, dtype=torch.int64, non_blocking=True
+            )
+        elif positions.dtype != torch.int64:
+            positions = positions.to(torch.int64)
+
+        max_position = (
+            self.max_position_hint
+            if self.max_position_hint is not None
+            else int(positions.max().item())
+        )
+        cos_sin_cache = self._ensure_rope_cache(max_position)
+
+        self._ensure_workspace(total_ctx, ctx_hidden.dtype)
+        assert self._proj_workspace is not None
+        assert self._k_workspace is not None
+        assert self._v_workspace is not None
+
+        proj_out_2d = self._proj_workspace[:total_ctx]
+        if self._mm_out_supported:
+            try:
+                torch.mm(ctx_hidden, self.flat_kv_weight_t, out=proj_out_2d)
+            except Exception:
+                self._mm_out_supported = False
+                proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t)
+        else:
+            proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t)
+
+        proj_out = proj_out_2d.view(total_ctx, self.n_layers, self.layer_out_dim)
+        tmp_k = self._k_workspace[:, :total_ctx]
+        tmp_v = self._v_workspace[:, :total_ctx]
+        cache_k, cache_v = _fused_norm_rope_stacked(
+            proj_out,
+            self.k_norm_weights,
+            self.eps_values,
+            cos_sin_cache,
+            positions,
+            self.num_kv_heads,
+            self.head_dim,
+            self.rotary_dim,
+            k_out=tmp_k,
+            v_out=tmp_v,
+        )
+        for layer_idx in range(self.n_layers):
+            write_layer_kv(layer_idx, cache_k[layer_idx], cache_v[layer_idx])

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

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _gather_rows_kernel(
+    idx_ptr,
+    s0,
+    d0,
+    n0,
+    s1,
+    d1,
+    n1,
+    s2,
+    d2,
+    n2,
+    s3,
+    d3,
+    n3,
+    HAS3: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    # One program == one (output row, column block). All buffers share the
+    # same gather index, so a single launch copies every buffer's row and
+    # the per-kernel launch bubbles between the old separate gathers vanish.
+    row = tl.program_id(0)
+    cb = tl.program_id(1)
+    src = tl.load(idx_ptr + row).to(tl.int64)
+    cols = cb * BLOCK + tl.arange(0, BLOCK)
+
+    m0 = cols < n0
+    tl.store(d0 + row * n0 + cols, tl.load(s0 + src * n0 + cols, mask=m0), mask=m0)
+
+    m1 = cols < n1
+    tl.store(d1 + row * n1 + cols, tl.load(s1 + src * n1 + cols, mask=m1), mask=m1)
+
+    m2 = cols < n2
+    tl.store(d2 + row * n2 + cols, tl.load(s2 + src * n2 + cols, mask=m2), mask=m2)
+
+    if HAS3:
+        m3 = cols < n3
+        tl.store(d3 + row * n3 + cols, tl.load(s3 + src * n3 + cols, mask=m3), mask=m3)
+
+
+def _row_width(buf: torch.Tensor) -> int:
+    """Flattened per-row element count (trailing dims), 1 for a 1-D buffer."""
+    return buf[0].numel() if buf.dim() > 1 else 1
+
+
+def _empty_like_rows(buf: torch.Tensor, m: int) -> torch.Tensor:
+    """Output buffer for `m` gathered rows of `buf` (same trailing dims/dtype/device)."""
+    return torch.empty((m, *buf.shape[1:]), dtype=buf.dtype, device=buf.device)
+
+
+def gather_spec_extras(
+    indices: torch.Tensor,
+    topk_p_buf: torch.Tensor,
+    topk_index_buf: torch.Tensor,
+    output_tokens_buf: torch.Tensor,
+    hidden_states_buf: Optional[torch.Tensor],
+):
+    """Gather spec extras (topk_p / topk_index / bonus_tokens / optional hidden
+    states) by a shared row index in a single fused Triton launch (one kernel
+    for all buffers) instead of one advanced-index gather per buffer.
+    `hidden_states_buf` is None when the build does not capture hidden states."""
+    # Source buffers are allocated once (torch.empty/full) and only ever mutated
+    # in place, so they are guaranteed row-contiguous. `indices` flows from
+    # several producers (req_pool_indices, filtered/merged future_indices); the
+    # kernel addresses it linearly, so normalize layout here (no-op when already
+    # contiguous) to avoid a silent wrong-result on a strided index tensor.
+    indices = indices.contiguous()
+    m = indices.shape[0]
+    has_hidden = hidden_states_buf is not None
+
+    topk_p = _empty_like_rows(topk_p_buf, m)
+    topk_index = _empty_like_rows(topk_index_buf, m)
+    bonus_tokens = _empty_like_rows(output_tokens_buf, m)
+    hidden_states = _empty_like_rows(hidden_states_buf, m) if has_hidden else None
+    if m == 0:
+        return topk_p, topk_index, bonus_tokens, hidden_states
+
+    n0 = _row_width(topk_p_buf)
+    n1 = _row_width(topk_index_buf)
+    n2 = _row_width(output_tokens_buf)
+    n3 = _row_width(hidden_states_buf) if has_hidden else 1
+    max_n = max(n0, n1, n2, n3)
+
+    # Dummy operands for the disabled hidden-states slot: the pointers must be
+    # valid even though the kernel never dereferences them (gated off by HAS3).
+    s3 = hidden_states_buf if has_hidden else indices
+    d3 = hidden_states if has_hidden else indices
+
+    block = min(1024, triton.next_power_of_2(max_n))
+    grid = (m, triton.cdiv(max_n, block))
+    _gather_rows_kernel[grid](
+        indices,
+        topk_p_buf,
+        topk_p,
+        n0,
+        topk_index_buf,
+        topk_index,
+        n1,
+        output_tokens_buf,
+        bonus_tokens,
+        n2,
+        s3,
+        d3,
+        n3,
+        HAS3=has_hidden,
+        BLOCK=block,
+    )
+    return topk_p, topk_index, bonus_tokens, hidden_states

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

@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cpu, is_npu
+
+_is_cpu = is_cpu()
+_is_npu = is_npu()
+
+if _is_cpu:
+    from sgl_kernel import rotate_input_ids_cpu
+
+
+@triton.jit
+def rotate_input_ids_kernel(
+    input_ids_ptr,
+    extend_start_loc_ptr,
+    extend_seq_lens_ptr,
+    topk_index_ptr,
+    select_index_ptr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    start_loc = tl.load(extend_start_loc_ptr + pid)
+    seq_len = tl.load(extend_seq_lens_ptr + pid)
+    new_token = tl.load(topk_index_ptr + pid)
+
+    num_elements_to_shift = seq_len - 1
+
+    for off in range(0, num_elements_to_shift, BLOCK_SIZE):
+        offsets = off + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < num_elements_to_shift
+
+        read_ptr = input_ids_ptr + start_loc + offsets + 1
+        val = tl.load(read_ptr, mask=mask)
+        tl.debug_barrier()
+
+        write_ptr = input_ids_ptr + start_loc + offsets
+        tl.store(write_ptr, val, mask=mask)
+        tl.debug_barrier()
+
+    if seq_len > 0:
+        if select_index_ptr is not None:
+            last_pos_ptr = input_ids_ptr + tl.load(select_index_ptr + pid)
+        else:
+            last_pos_ptr = input_ids_ptr + start_loc + seq_len - 1
+        tl.store(last_pos_ptr, new_token)
+
+
+def rotate_input_ids(
+    input_ids, extend_start_loc, extend_seq_lens, topk_index, select_index=None
+):
+    if _is_cpu:
+        rotate_input_ids_cpu(
+            input_ids,
+            extend_start_loc,
+            extend_seq_lens,
+            topk_index,
+            select_index,
+        )
+        return input_ids
+
+    batch_size = extend_seq_lens.shape[0]
+
+    # rotate_input_ids_triton skipped: batch_size=0 (empty extend_seq_lens).
+    # This is expected when a DP rank has no requests.
+    # TODO: @iforgetmyname Remove NPU-specific guard after triton-ascend fixes zero-sized grid kernel launch abort
+    if batch_size == 0 and _is_npu:
+        return input_ids
+
+    BLOCK_SIZE = 4096 if select_index is not None else 8
+    grid = (batch_size,)
+
+    rotate_input_ids_kernel[grid](
+        input_ids,
+        extend_start_loc,
+        extend_seq_lens,
+        topk_index,
+        select_index,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return input_ids

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

@@ -0,0 +1,281 @@
+# Copyright 2023-2026 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.
+# ==============================================================================
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def sgl_build_tree_kernel_efficient_triton(
+    parent_list_ptr,
+    selected_index_ptr,
+    verified_seq_len_ptr,
+    seq_len_prefix_sum_ptr,
+    tree_mask_ptr,
+    positions_ptr,
+    retrieve_index_ptr,
+    retrieve_next_token_ptr,
+    retrieve_next_sibling_ptr,
+    topk: tl.constexpr,
+    depth: tl.constexpr,
+    draft_token_num: tl.constexpr,
+    tree_mask_mode: tl.constexpr,
+    batch_size: tl.constexpr,
+    parent_list_stride: tl.constexpr,
+    selected_index_stride: tl.constexpr,
+):
+    """
+    Triton kernel for building EAGLE tree structure.
+    Each program handles one batch item (batch_idx).
+    """
+    batch_idx = tl.program_id(0)
+
+    # Calculate seq_tree_idx
+    seq_len = tl.load(verified_seq_len_ptr + batch_idx)
+    seq_len_prefix_sum = tl.load(seq_len_prefix_sum_ptr + batch_idx)
+
+    # Cast initial value to match the dtype of loaded tensors to avoid type inconsistency
+    seq_tree_idx = (
+        tl.cast(draft_token_num * draft_token_num * batch_idx, seq_len.dtype)
+        + seq_len_prefix_sum * draft_token_num
+    )
+
+    positions_offset = batch_idx * draft_token_num
+    tl.store(positions_ptr + positions_offset, seq_len)
+
+    retrieve_index_offset = batch_idx * draft_token_num
+
+    # Build retrieval index structure (reverse loop from draft_token_num-1 to 1)
+    for i in range(draft_token_num - 1, 0, -1):
+        current_token_idx = retrieve_index_offset + i
+        tl.store(
+            retrieve_index_ptr + batch_idx * draft_token_num + i,
+            current_token_idx,
+        )
+
+        parent_tb_idx = (
+            tl.load(selected_index_ptr + batch_idx * selected_index_stride + (i - 1))
+            // topk
+        )
+        parent_position = 0
+        found = 0
+
+        if parent_tb_idx == 0:
+            found = 1
+        else:
+            parent_token_idx = tl.load(
+                parent_list_ptr + batch_idx * parent_list_stride + parent_tb_idx
+            )
+
+            # Find parent position
+            for pp in range(draft_token_num - 1):
+                if found == 0:
+                    sel_idx = tl.load(
+                        selected_index_ptr + batch_idx * selected_index_stride + pp
+                    )
+                    if sel_idx == parent_token_idx:
+                        parent_position = pp + 1
+                        found = 1
+
+        if found == 1:
+            # Update next token links
+            next_tok_addr = (
+                retrieve_next_token_ptr + batch_idx * draft_token_num + parent_position
+            )
+            next_tok = tl.load(next_tok_addr)
+
+            if next_tok == -1:
+                tl.store(next_tok_addr, i)
+            else:
+                tl.store(next_tok_addr, i)
+                tl.store(
+                    retrieve_next_sibling_ptr + batch_idx * draft_token_num + i,
+                    next_tok,
+                )
+
+    tl.store(retrieve_index_ptr + batch_idx * draft_token_num, retrieve_index_offset)
+
+    # Process all draft token indices for tree mask
+    for draft_tokenx in range(draft_token_num):
+        if tree_mask_mode == 0:  # FULL_MASK
+            token_tree_idx = (
+                seq_tree_idx + (seq_len + draft_token_num) * draft_tokenx + seq_len + 1
+            )
+        else:
+            token_tree_idx = (
+                draft_token_num * draft_token_num * batch_idx
+                + draft_token_num * draft_tokenx
+                + 1
+            )
+
+        tl.store(tree_mask_ptr + token_tree_idx - 1, 1)
+        for i in range(draft_token_num - 1):
+            tl.store(tree_mask_ptr + token_tree_idx + i, 0)
+
+        if draft_tokenx > 0:
+            # Build tree path for draft_tokenx > 0
+            cur_position = draft_tokenx - 1
+            position = 0
+            should_continue = 1
+
+            for _ in range(depth):
+                if should_continue:
+                    position += 1
+                    tl.store(tree_mask_ptr + token_tree_idx + cur_position, 1)
+
+                    parent_tb_idx = (
+                        tl.load(
+                            selected_index_ptr
+                            + batch_idx * selected_index_stride
+                            + cur_position
+                        )
+                        // topk
+                    )
+                    if parent_tb_idx == 0:
+                        should_continue = 0
+                    else:
+                        parent_token_idx = tl.load(
+                            parent_list_ptr
+                            + batch_idx * parent_list_stride
+                            + parent_tb_idx
+                        )
+
+                        # Find cur_position for next iteration
+                        found = 0
+                        for cp in range(draft_token_num - 1):
+                            if found == 0:
+                                if (
+                                    tl.load(
+                                        selected_index_ptr
+                                        + batch_idx * selected_index_stride
+                                        + cp
+                                    )
+                                    == parent_token_idx
+                                ):
+                                    cur_position = cp
+                                    found = 1
+                        if found == 0:
+                            should_continue = 0
+
+            tl.store(
+                positions_ptr + batch_idx * draft_token_num + draft_tokenx,
+                position + seq_len,
+            )
+
+
+@triton.jit
+def verify_tree_greedy_kernel_triton(
+    predicts_ptr,
+    accept_index_ptr,
+    accept_token_num_ptr,
+    candidates_ptr,
+    retrieve_index_ptr,
+    retrieve_next_token_ptr,
+    retrieve_next_sibling_ptr,
+    target_predict_ptr,
+    batch_size: tl.constexpr,
+    num_speculative_tokens: tl.constexpr,
+    num_draft_tokens: tl.constexpr,
+):
+    """
+    Triton kernel for verifying EAGLE tree in greedy mode.
+    Each program handles one batch item.
+    """
+    bx = tl.program_id(0)
+
+    # Initialize
+    last_accept_retrieve_idx = tl.load(retrieve_index_ptr + bx * num_draft_tokens)
+    tl.store(accept_index_ptr + bx * num_speculative_tokens, last_accept_retrieve_idx)
+    # Cast to match dtype of loaded tensors to avoid type inconsistency
+    num_accept_tokens = tl.cast(0, last_accept_retrieve_idx.dtype)
+    cur_index = tl.cast(0, last_accept_retrieve_idx.dtype)
+
+    # Tree traversal loop
+    should_continue = 1
+    for j in range(1, num_speculative_tokens):
+        if should_continue:  # Early exit guard
+            cur_index = tl.load(
+                retrieve_next_token_ptr + bx * num_draft_tokens + cur_index
+            )
+
+            # Load target token once per level (before sibling search)
+            # last_accept_retrieve_idx is constant during sibling traversal
+            target_row = last_accept_retrieve_idx // num_draft_tokens
+            target_col = last_accept_retrieve_idx % num_draft_tokens
+            target_token = tl.load(
+                target_predict_ptr + target_row * num_draft_tokens + target_col
+            )
+
+            # Traverse siblings
+            found_match = 0
+            for _ in range(num_draft_tokens):  # Max iterations = num_draft_tokens
+                if found_match == 0:  # Early exit guard
+                    # Check if we've reached end of sibling list
+                    is_valid = cur_index != -1
+
+                    # Use masked loads with safe address (0 when invalid)
+                    safe_cur_index = (
+                        cur_index * is_valid
+                    )  # 0 if invalid, cur_index if valid
+                    safe_index = bx * num_draft_tokens + safe_cur_index
+
+                    # Load draft token info (loads from index 0 when invalid, but we won't use it)
+                    draft_index = tl.load(retrieve_index_ptr + safe_index)
+                    draft_token = tl.load(candidates_ptr + safe_index)
+
+                    # Check for token match (only valid when is_valid is True)
+                    token_match = is_valid & (draft_token == target_token)
+
+                    # Accept token using predicated stores (only write if matched)
+                    tl.store(
+                        predicts_ptr + last_accept_retrieve_idx,
+                        target_token,
+                        mask=token_match,
+                    )
+                    next_num_accept_tokens = num_accept_tokens + 1
+                    tl.store(
+                        accept_index_ptr
+                        + bx * num_speculative_tokens
+                        + next_num_accept_tokens,
+                        draft_index,
+                        mask=token_match,
+                    )
+
+                    num_accept_tokens = num_accept_tokens + token_match
+                    last_accept_retrieve_idx = (
+                        token_match * draft_index
+                        + (~token_match) * last_accept_retrieve_idx
+                    )
+                    found_match = token_match * 1 + (~is_valid) * (-1)
+
+                    # Masked load: only load next sibling when no match (hardware predication)
+                    # When matched: returns cur_index (other); when not matched: loads sibling
+                    cur_index = tl.load(
+                        retrieve_next_sibling_ptr + safe_index,
+                        mask=~token_match
+                        & is_valid,  # Only load when valid and NOT matched
+                        other=cur_index,  # Keep cur_index when matched or invalid
+                    )
+
+            if found_match != 1:
+                should_continue = 0
+
+    # Store final results
+    tl.store(accept_token_num_ptr + bx, num_accept_tokens)
+
+    target_row = last_accept_retrieve_idx // num_draft_tokens
+    target_col = last_accept_retrieve_idx % num_draft_tokens
+    final_target = tl.load(
+        target_predict_ptr + target_row * num_draft_tokens + target_col
+    )
+    tl.store(predicts_ptr + last_accept_retrieve_idx, final_target)

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 ``"<group>.<name>"`` 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)

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, ``"<group>.<name>"``.
+    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()

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, ``"<group>.<name>"`` (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