chore: import upstream snapshot with attribution

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