minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Kernel] Phase 3+4: move JIT infra + operator groups into sglang.kernels (RFC #29630) (#31666)

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-22 11:28:27 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent d708969f68
commit 246b3c3eaf
155 changed files with 1293 additions and 1229 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/__main__.py
+4 -4
View File
@@ -7,10 +7,10 @@ import subprocess
from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
from sglang.jit_kernel.utils import get_jit_cuda_arch, override_jit_cuda_arch
from sglang.jit_kernel.utils.arch import get_default_target_flags
from sglang.jit_kernel.utils.compile import DEFAULT_INCLUDE
from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES
from sglang.kernels.jit.utils import get_jit_cuda_arch, override_jit_cuda_arch
from sglang.kernels.jit.utils.arch import get_default_target_flags
from sglang.kernels.jit.utils.compile import DEFAULT_INCLUDE
from sglang.kernels.jit.utils.deps import REGISTERED_DEPENDENCIES
def _clangd_major_version() -> int | None:
python/sglang/jit_kernel/activation.py
+3 -166
View File
@@ -1,168 +1,5 @@
from __future__ import annotations
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.activation._jit_activation."""
from typing import TYPE_CHECKING, Optional
from sglang.kernels.ops.activation import _jit_activation as _impl
import torch
from sglang.jit_kernel.utils import (
cache_once,
get_jit_cuda_arch,
is_arch_support_pdl,
is_hip_runtime,
load_jit,
make_cpp_args,
)
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
def _fast_math_flags() -> list[str]:
# Mirrors sgl-kernel's CMake policy: fast-math on SM90, precise on
# SM100+ (Blackwell needs bit-exact expf), off on HIP (clang rejects).
if is_hip_runtime():
return []
if get_jit_cuda_arch().major >= 10:
return []
return ["--use_fast_math"]
@cache_once
def _jit_activation_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype, is_arch_support_pdl())
return load_jit(
"activation",
*args,
cuda_files=["elementwise/activation.cuh"],
extra_cuda_cflags=_fast_math_flags(),
cuda_wrappers=[
("run_activation", f"ActivationKernel<{args}>::run_activation"),
(
"run_activation_filtered",
f"ActivationKernel<{args}>::run_activation_filtered",
),
(
"run_unary_activation",
f"ActivationKernel<{args}>::run_unary_activation",
),
],
)
SUPPORTED_ACTIVATIONS = {"silu", "gelu", "gelu_tanh"}
SUPPORTED_UNARY_ACTIVATIONS = {"relu2"}
@register_custom_op(mutates_args=["out"])
def _run_activation_inplace(
op_name: str, input: torch.Tensor, out: torch.Tensor
) -> None:
hidden_size = input.shape[-1] // 2
module = _jit_activation_module(input.dtype)
input_2d = input.view(-1, hidden_size * 2)
out_2d = out.view(-1, hidden_size)
module.run_activation(input_2d, out_2d, op_name)
@register_custom_op(mutates_args=["out"])
def _run_activation_filtered_inplace(
op_name: str,
input: torch.Tensor,
out: torch.Tensor,
expert_ids: torch.Tensor,
expert_step: int,
) -> None:
hidden_size = input.shape[-1] // 2
module = _jit_activation_module(input.dtype)
input_2d = input.view(-1, hidden_size * 2)
out_2d = out.view(-1, hidden_size)
module.run_activation_filtered(input_2d, out_2d, expert_ids, expert_step, op_name)
def run_activation(
op_name: str,
input: torch.Tensor,
out: Optional[torch.Tensor],
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
"""Apply ``op_name`` activation followed by element-wise multiplication.
When ``expert_ids`` is provided, output rows are skipped for tokens whose
routed expert id is ``-1``. ``expert_step`` is 1 for per-token routing and
``BLOCK_SIZE_M`` for sorted/TMA routing — i.e. ``expert_ids[token_id //
expert_step]`` is consulted before computing each row.
"""
assert op_name in SUPPORTED_ACTIVATIONS, f"Unsupported activation: {op_name}"
hidden_size = input.shape[-1] // 2
if out is None:
out = input.new_empty(*input.shape[:-1], hidden_size)
if expert_ids is None:
_run_activation_inplace(op_name, input, out)
else:
_run_activation_filtered_inplace(op_name, input, out, expert_ids, expert_step)
return out
@register_custom_op(mutates_args=["out"])
def _run_unary_activation_inplace(
op_name: str, input: torch.Tensor, out: torch.Tensor
) -> None:
last = input.shape[-1]
module = _jit_activation_module(input.dtype)
module.run_unary_activation(input.view(-1, last), out.view(-1, last), op_name)
def run_unary_activation(
op_name: str,
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Apply a standalone (non-gated) element-wise activation: ``out = act(input)``.
Unlike :func:`run_activation`, there is no gate/up split — ``input`` and
``out`` share the same shape.
"""
assert (
op_name in SUPPORTED_UNARY_ACTIVATIONS
), f"Unsupported unary activation: {op_name}"
if out is None:
out = torch.empty_like(input)
_run_unary_activation_inplace(op_name, input, out)
return out
def relu2(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Squared ReLU: ``out = max(0, input) ** 2`` (element-wise)."""
return run_unary_activation("relu2", input, out)
def silu_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("silu", input, out, expert_ids, expert_step)
def gelu_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("gelu", input, out, expert_ids, expert_step)
def gelu_tanh_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("gelu_tanh", input, out, expert_ids, expert_step)
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/add_constant.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/all_reduce.py
+2 -2
View File
@@ -7,14 +7,14 @@ import torch
import tvm_ffi
from tvm_ffi import Module
from sglang.jit_kernel.utils import (
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
lazy_register_class,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
class AllReduceAlgo(enum.Enum):
python/sglang/jit_kernel/awq_dequantize.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/awq_marlin_repack.py
+1 -1
View File
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/benchmark/marker.py
+1 -1
View File
@@ -21,7 +21,7 @@ from typing import (
import torch
from sglang.jit_kernel.utils import cache_once
from sglang.kernels.jit.utils import cache_once
from sglang.utils import is_in_ci
F = TypeVar("F", bound=Callable[..., "BenchResult"])
python/sglang/jit_kernel/clamp_position.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/concat_mla.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/diffusion/norm_scale_shift_native.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/diffusion/qknorm_rope.py
+1 -1
View File
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/diffusion/residual_gate_add.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/dsv32/elementwise.py
+1 -1
View File
@@ -2,7 +2,7 @@
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/dsv3_fused_a_gemm.py
+3 -88
View File
@@ -1,90 +1,5 @@
"""
JIT kernel for DeepSeek V3 fused QKV-A GEMM (min-latency).
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm."""
Runtime-compiled CUDA C++ kernel for SM90+ (Hopper) GPUs.
Shapes: hd_in a multiple of 256, hd_out a multiple of 16, num_tokens 1-16, bfloat16.
"""
from sglang.kernels.ops.gemm import _jit_dsv3_fused_a_gemm as _impl
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.common import direct_register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_dsv3_fused_a_gemm_module(hd_in: int, hd_out: int, use_pdl: bool) -> Module:
args = make_cpp_args(hd_in, hd_out, use_pdl)
return load_jit(
"dsv3_fused_a_gemm",
*args,
cuda_files=["gemm/dsv3_fused_a_gemm.cuh"],
cuda_wrappers=[
("dsv3_fused_a_gemm", f"DSV3FusedAGemmKernel<{args}>::run"),
],
)
def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
assert mat_a.stride(1) == 1, "mat_a must be row-major [M, K]"
output = torch.empty(
(mat_a.shape[0], mat_b.shape[1]),
device=mat_a.device,
dtype=mat_a.dtype,
)
module = _jit_dsv3_fused_a_gemm_module(
mat_a.shape[1], mat_b.shape[1], is_arch_support_pdl()
)
module.dsv3_fused_a_gemm(mat_a, mat_b, output)
return output
def _dsv3_fused_a_gemm_fake(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
return mat_a.new_empty((mat_a.shape[0], mat_b.shape[1]), dtype=torch.bfloat16)
direct_register_custom_op(
op_name="jit_dsv3_fused_a_gemm",
op_func=_dsv3_fused_a_gemm_run,
mutates_args=[],
fake_impl=_dsv3_fused_a_gemm_fake,
)
@debug_kernel_api
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 kernel (JIT variant).
Args:
mat_a: Input tensor of shape [num_tokens, hd_in], bfloat16, row-major.
hd_in must be a multiple of 256 and num_tokens in [1, 16].
mat_b: Weight tensor of shape [hd_in, hd_out], bfloat16, column-major
(i.e. ``weight.T`` of a row-major [hd_out, hd_in] weight).
hd_out must be a multiple of 16.
output: Optional pre-allocated output tensor of shape [num_tokens, hd_out].
Returns:
Output tensor of shape [num_tokens, hd_out].
"""
result = torch.ops.sglang.jit_dsv3_fused_a_gemm(mat_a, mat_b)
if output is not None:
output.copy_(result)
return output
return result
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/dsv3_router_gemm.py
+3 -90
View File
@@ -1,92 +1,5 @@
"""
JIT kernel for DeepSeek V3 router GEMM.
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.gemm._jit_dsv3_router_gemm."""
Runtime-compiled CUDA C++ kernel for SM90+ (Hopper) GPUs.
Supports num_experts in {256, 384}, hidden_dim a multiple of 1024, num_tokens 1-16.
"""
from sglang.kernels.ops.gemm import _jit_dsv3_router_gemm as _impl
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_dsv3_router_gemm_module(
num_experts: int,
hidden_dim: int,
use_pdl: bool,
out_float: bool,
) -> Module:
args = make_cpp_args(num_experts, hidden_dim, use_pdl, out_float)
return load_jit(
"dsv3_router_gemm",
*args,
cuda_files=["gemm/dsv3_router_gemm.cuh"],
cuda_wrappers=[
("dsv3_router_gemm", f"DSV3RouterGemmKernel<{args}>::run"),
],
)
@register_custom_op(
op_name="dsv3_router_gemm",
mutates_args=["output"],
)
def _dsv3_router_gemm_custom_op(
hidden_states: torch.Tensor,
router_weights: torch.Tensor,
output: torch.Tensor,
) -> None:
num_experts = router_weights.shape[0]
hidden_dim = hidden_states.shape[1]
out_float = output.dtype == torch.float32
module = _jit_dsv3_router_gemm_module(
num_experts, hidden_dim, is_arch_support_pdl(), out_float
)
module.dsv3_router_gemm(hidden_states, router_weights, output)
return None
@debug_kernel_api
def dsv3_router_gemm(
hidden_states: torch.Tensor,
router_weights: torch.Tensor,
out_dtype: torch.dtype = torch.bfloat16,
output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
DeepSeek V3 router GEMM kernel (JIT variant).
Args:
hidden_states: Input tensor of shape [num_tokens, hidden_dim], bfloat16.
hidden_dim must be a multiple of 1024 and num_tokens in [1, 16].
router_weights: Weight tensor of shape [num_experts, hidden_dim], bfloat16.
out_dtype: Output dtype, either torch.bfloat16 or torch.float32.
output: Optional pre-allocated output tensor.
Returns:
Output tensor of shape [num_tokens, num_experts].
"""
if output is None:
output = torch.empty(
hidden_states.shape[0],
router_weights.shape[0],
device=hidden_states.device,
dtype=out_dtype,
)
_dsv3_router_gemm_custom_op(hidden_states, router_weights, output)
return output
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/dsv4/attn.py
+1 -1
View File
@@ -4,7 +4,7 @@ import torch
import triton
import triton.language as tl
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
is_hip_runtime,
python/sglang/jit_kernel/dsv4/compress.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/dsv4/compress_old.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/dsv4/elementwise.py
+1 -1
View File
@@ -2,7 +2,7 @@ from typing import Optional, Tuple
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/dsv4/fp8_wo_a.py
+2 -2
View File
@@ -4,13 +4,13 @@ from typing import TYPE_CHECKING, Tuple
import torch
from sglang.jit_kernel.utils import (
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.custom_op import register_custom_op
from .utils import make_name
python/sglang/jit_kernel/dsv4/moe.py
+1 -1
View File
@@ -2,7 +2,7 @@ from typing import Optional, Tuple
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
is_hip_runtime,
python/sglang/jit_kernel/dsv4/online_c128_mtp.py
+1 -1
View File
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, List, Optional
import torch
from sglang.jit_kernel.dsv4.utils import make_name
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.environ import envs
if TYPE_CHECKING:
python/sglang/jit_kernel/dsv4/topk.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import Optional
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
is_hip_runtime,
python/sglang/jit_kernel/fixup_zero_kv.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/flash_attention_v3.py
+1 -1
View File
@@ -4,8 +4,8 @@ from typing import Optional, Union
import torch
from sglang.jit_kernel.utils import cache_once
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once
from sglang.srt.environ import envs
from sglang.srt.utils import get_device_capability, is_musa
python/sglang/jit_kernel/flash_attn/cute/interface.py
+1 -1
View File
@@ -15,7 +15,7 @@ from quack.compile_utils import make_fake_tensor as fake_tensor
from sglang.jit_kernel.flash_attn.cute.cache_utils import get_jit_cache
from sglang.jit_kernel.flash_attn.cute.testing import is_fake_mode
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.kernels.jit.utils import is_arch_support_pdl
if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None:
from sglang.jit_kernel.flash_attn.cute import cute_dsl_ptxas # noqa: F401
python/sglang/jit_kernel/fp8_blockwise_gemm.py
+1 -1
View File
@@ -5,8 +5,8 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit, override_jit_cuda_arch
from sglang.srt.utils.common import is_sm120_supported
from sglang.srt.utils.custom_op import register_custom_op
python/sglang/jit_kernel/fused_eh_norm.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/fused_fp8_qkv_kv_cache.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/fused_metadata_copy.py
+1 -1
View File
@@ -15,7 +15,7 @@ from typing import Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
logger = logging.getLogger(__name__)
python/sglang/jit_kernel/fused_qknorm_rope.py
+1 -1
View File
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/fused_store_index_cache.py
+2 -2
View File
@@ -13,13 +13,13 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/gptq_marlin.py
+1 -1
View File
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from sgl_kernel.scalar_type import ScalarType
python/sglang/jit_kernel/gptq_marlin_repack.py
+1 -1
View File
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/hadamard.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Callable
import torch
from sglang.jit_kernel.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/hicache.py
+1 -1
View File
@@ -3,8 +3,8 @@ from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
import torch
python/sglang/jit_kernel/hisparse.py
+1 -1
View File
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import load_jit, make_cpp_args
from sglang.kernels.jit.utils import load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/inkling_all_reduce.py
+1 -1
View File
@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, empty_sentinel, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, empty_sentinel, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/inkling_ar_fused.py
+1 -1
View File
@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, empty_sentinel, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, empty_sentinel, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/inkling_ar_scattered_sconv.py
+1 -1
View File
@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/inkling_attn_prologue.py
+1 -1
View File
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
empty_sentinel,
is_arch_support_pdl,
python/sglang/jit_kernel/inkling_gate_topk_renorm.py
+1 -1
View File
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/inkling_rel_proj.py
+1 -1
View File
@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
empty_sentinel,
is_arch_support_pdl,
python/sglang/jit_kernel/inkling_row_scale.py
+1 -1
View File
@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/inkling_sconv.py
+1 -1
View File
@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/kpool_topk_transform.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/kv_canary/verify.py
+1 -1
View File
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Final
import torch
from sglang.jit_kernel.kv_canary import consts
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/kv_canary/write.py
+1 -1
View File
@@ -11,7 +11,7 @@ from sglang.jit_kernel.kv_canary.verify import (
_assert_contiguous,
_build_real_kv_source_abi,
)
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/kvcache.py
+1 -1
View File
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/lplb/cuda_solver.py
+1 -1
View File
@@ -17,7 +17,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
get_jit_cuda_arch,
load_jit,
python/sglang/jit_kernel/minimax_decode_topk.py
+1 -1
View File
@@ -17,7 +17,7 @@ from typing import TYPE_CHECKING, Tuple
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/minimax_qknorm_rope.py
+1 -1
View File
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, List, Sequence, Tuple
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/minimax_quant_ue8m0.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Tuple
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/minimax_store_kv_index.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/mla_kv_pack_quantize_fp8.py
+1 -1
View File
@@ -11,7 +11,7 @@ import torch
import triton
import triton.language as tl
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.kernels.jit.utils import is_arch_support_pdl
@triton.jit
python/sglang/jit_kernel/moe_align.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/moe_finalize_fuse_shared.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
@cache_once
python/sglang/jit_kernel/moe_fused_gate.py
+1 -1
View File
@@ -7,8 +7,8 @@ import torch
import triton
import triton.language as tl
from sglang.jit_kernel.utils import cache_once, is_arch_support_pdl, load_jit
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/moe_lora_align.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/moe_permute_prepare.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Tuple
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/moe_topk_sigmoid.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/moe_wna16_marlin.py
+1 -1
View File
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from sgl_kernel.scalar_type import ScalarType
python/sglang/jit_kernel/mxfp8.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
load_jit,
make_cpp_args,
python/sglang/jit_kernel/ngram_corpus.py
+1 -1
View File
@@ -7,7 +7,7 @@ import numpy as np
import torch
import tvm_ffi
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
_MATCH_TYPE_MAP = {"BFS": 0, "PROB": 1}
python/sglang/jit_kernel/ngram_embedding.py
+1 -1
View File
@@ -2,8 +2,8 @@ from __future__ import annotations
from typing import TYPE_CHECKING
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
import torch
python/sglang/jit_kernel/norm.py
+3 -177
View File
@@ -1,179 +1,5 @@
from __future__ import annotations
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.layernorm._jit_norm."""
import logging
from typing import TYPE_CHECKING, Optional
from sglang.kernels.ops.layernorm import _jit_norm as _impl
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
if TYPE_CHECKING:
from tvm_ffi.module import Module
logger = logging.getLogger(__name__)
@cache_once
def _jit_qknorm_module(head_dim: int, dtype: torch.dtype) -> Module:
args = make_cpp_args(head_dim, is_arch_support_pdl(), dtype)
return load_jit(
"qknorm",
*args,
cuda_files=["elementwise/qknorm.cuh"],
cuda_wrappers=[("qknorm", f"QKNormKernel<{args}>::run")],
)
_RMSNORM_WARP_SIZES = frozenset({64, 128, 256})
_RMSNORM_MAX_HIDDEN_SIZE = 16384
_RMSNORM_HALF_BLOCK_MIN_SIZE = 2048
def _is_supported_rmsnorm_hidden_size(d: int) -> bool:
return d in _RMSNORM_WARP_SIZES or (
(d > 256 and d % 256 == 0 and d <= 8192)
or (d >= 8192 and d % 512 == 0 and d <= 16384)
)
def _rmsnorm_kernel_class(hidden_size: int) -> str:
if hidden_size in _RMSNORM_WARP_SIZES:
return "RMSNormWarpKernel"
if hidden_size == 512:
return "RMSNormHalfKernel"
if hidden_size >= _RMSNORM_HALF_BLOCK_MIN_SIZE:
if hidden_size % 512 == 0:
return "RMSNormHalfKernel"
return "RMSNormKernel"
@cache_once
def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module:
args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype)
kernel_class = f"{_rmsnorm_kernel_class(hidden_size)}<{args}>"
return load_jit(
"rmsnorm",
*args,
cuda_files=["elementwise/rmsnorm.cuh"],
cuda_wrappers=[("rmsnorm", f"{kernel_class}::run")],
)
def is_supported_jit_fused_add_rmsnorm_hidden_size(hidden_size: int) -> bool:
return hidden_size > 0 and hidden_size % 16 == 0 and hidden_size <= 8192
@cache_once
def _jit_fused_add_rmsnorm_module(
dtype: torch.dtype, cast_x_before_out_mul: bool
) -> Module:
args = make_cpp_args(cast_x_before_out_mul, dtype)
return load_jit(
"fused_add_rmsnorm",
*args,
cuda_files=["elementwise/fused_add_rmsnorm.cuh"],
cuda_wrappers=[("fused_add_rmsnorm", f"FusedAddRMSNormKernel<{args}>::run")],
)
@cache_once
def _jit_qknorm_across_heads_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype)
return load_jit(
"qknorm_across_heads",
*args,
cuda_files=["elementwise/qknorm_across_heads.cuh"],
cuda_wrappers=[
("qknorm_across_heads", f"QKNormAcrossHeadsKernel<{args}>::run")
],
)
@torch.compiler.assume_constant_result
@cache_once
def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool:
if head_dim not in [64, 128, 256, 512, 1024]:
logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel")
return False
try:
_jit_qknorm_module(head_dim, dtype)
return True
except Exception as e:
logger.warning(f"Failed to load JIT QK-Norm kernel: {e}")
return False
@debug_kernel_api
def fused_inplace_qknorm(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
eps: float = 1e-6,
*,
head_dim: int = 0,
) -> None:
head_dim = head_dim or q.size(-1)
module = _jit_qknorm_module(head_dim, q.dtype)
module.qknorm(q, k, q_weight, k_weight, eps)
@debug_kernel_api
def rmsnorm(
input: torch.Tensor,
weight: torch.Tensor,
out: Optional[torch.Tensor] = None,
eps: float = 1e-6,
) -> None:
out = out if out is not None else input
hidden_size = input.size(-1)
if not _is_supported_rmsnorm_hidden_size(hidden_size):
raise RuntimeError(
f"jit rmsnorm: unsupported hidden_size={hidden_size}. "
f"Supported: {sorted(_RMSNORM_WARP_SIZES)}, and multiples of 256 in "
f"(256, {_RMSNORM_MAX_HIDDEN_SIZE}]."
)
module = _jit_rmsnorm_module(hidden_size, input.dtype)
module.rmsnorm(input, weight, out, eps)
@debug_kernel_api
def fused_add_rmsnorm(
input: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
eps: float = 1e-6,
*,
cast_x_before_out_mul: bool = False,
) -> None:
module = _jit_fused_add_rmsnorm_module(input.dtype, cast_x_before_out_mul)
module.fused_add_rmsnorm(input, residual, weight, eps)
@debug_kernel_api
def fused_inplace_qknorm_across_heads(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
eps: float = 1e-6,
) -> None:
"""
Fused inplace QK normalization across all heads.
Args:
q: Query tensor of shape [batch_size, num_heads * head_dim]
k: Key tensor of shape [batch_size, num_heads * head_dim]
q_weight: Query weight tensor of shape [num_heads * head_dim]
k_weight: Key weight tensor of shape [num_heads * head_dim]
eps: Epsilon for numerical stability
"""
module = _jit_qknorm_across_heads_module(q.dtype)
module.qknorm_across_heads(q, k, q_weight, k_weight, eps)
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/per_tensor_quant_fp8.py
+3 -75
View File
@@ -1,77 +1,5 @@
from __future__ import annotations
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8."""
from typing import TYPE_CHECKING
from sglang.kernels.ops.quantization import _jit_per_tensor_quant_fp8 as _impl
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_per_tensor_quant_fp8_module(is_static: bool, dtype: torch.dtype) -> Module:
args = make_cpp_args(is_static, dtype)
return load_jit(
"per_tensor_quant_fp8",
*args,
cuda_files=["gemm/per_tensor_quant_fp8.cuh"],
cuda_wrappers=[("per_tensor_quant_fp8", f"per_tensor_quant_fp8<{args}>")],
)
@register_custom_op(
op_name="per_tensor_quant_fp8",
mutates_args=["output_q", "output_s"],
)
def per_tensor_quant_fp8(
input: torch.Tensor,
output_q: torch.Tensor,
output_s: torch.Tensor,
is_static: bool = False,
) -> None:
"""
Per-tensor quantization to FP8 format.
Args:
input: Input tensor to quantize (float, half, or bfloat16)
output_q: Output quantized tensor (fp8_e4m3)
output_s: Output scale tensor (float scalar or 1D tensor with 1 element)
is_static: If True, assumes scale is pre-computed and skips absmax computation
"""
module = _jit_per_tensor_quant_fp8_module(is_static, input.dtype)
module.per_tensor_quant_fp8(input.view(-1), output_q.view(-1), output_s.view(-1))
@cache_once
def _jit_per_tensor_absmax_fp8_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype)
return load_jit(
"per_tensor_absmax_fp8",
*args,
cuda_files=["gemm/per_tensor_quant_fp8.cuh"],
cuda_wrappers=[("per_tensor_absmax_fp8", f"per_tensor_absmax_fp8<{args}>")],
)
@register_custom_op(
op_name="per_tensor_absmax_fp8",
mutates_args=["output_s"],
)
def per_tensor_absmax_fp8(
input: torch.Tensor,
output_s: torch.Tensor,
) -> None:
"""Compute scale = max(abs(input)) / fp8_e4m3_max via atomic-max reduction.
The caller must zero-initialise ``output_s`` before the call (the kernel
uses ``atomic_max`` across blocks, so starting from 0 is required).
Args:
input: Input tensor (float16, bfloat16, or float32). Any shape.
output_s: Pre-allocated float32 tensor of shape (1,), zero-initialised.
"""
module = _jit_per_tensor_absmax_fp8_module(input.dtype)
module.per_tensor_absmax_fp8(input.view(-1), output_s.view(-1))
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/per_token_group_quant.py
+3 -221
View File
@@ -1,223 +1,5 @@
from __future__ import annotations
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.quantization._jit_per_token_group_quant."""
from typing import TYPE_CHECKING, Optional, Tuple
from sglang.kernels.ops.quantization import _jit_per_token_group_quant as _impl
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
_SUPPORTED_INPUT_DTYPES = (torch.bfloat16, torch.float16)
_SUPPORTED_OUTPUT_DTYPES = (torch.float8_e4m3fn, torch.int8)
_SUPPORTED_GROUP_SIZES = (16, 32, 64, 128, 256)
@cache_once
def _jit_module(
in_dtype: torch.dtype,
out_dtype: torch.dtype,
group_size: int,
scale_ue8m0: bool,
row_major: bool,
aligned: bool,
fuse_silu_and_mul: bool,
masked_layout: bool,
use_pdl: bool,
) -> Module:
assert in_dtype in _SUPPORTED_INPUT_DTYPES
assert out_dtype in _SUPPORTED_OUTPUT_DTYPES
assert group_size in _SUPPORTED_GROUP_SIZES
trait_args = make_cpp_args(
in_dtype,
out_dtype,
group_size,
scale_ue8m0,
row_major,
aligned,
fuse_silu_and_mul,
use_pdl,
)
launcher = (
"PerTokenGroupQuantMaskedKernel"
if masked_layout
else "PerTokenGroupQuantFlatKernel"
)
return load_jit(
"per_token_group_quant",
*trait_args,
"masked" if masked_layout else "flat",
cuda_files=["gemm/per_token_group_quant.cuh"],
cuda_wrappers=[("per_token_group_quant", f"{launcher}<{trait_args}>::run")],
extra_cuda_cflags=["--use_fast_math"],
)
def _infer_scale_layout(
output_s: torch.Tensor, scale_ue8m0: bool, num_groups: int
) -> Tuple[bool, bool]:
"""Return ``(row_major, aligned)`` for ``output_s``.
Column-major (transposed) scale buffers have token stride 1 and a larger
group stride; row-major buffers are contiguous.
"""
row_major = output_s.stride(-2) >= output_s.stride(-1)
if output_s.dtype == torch.int32:
if not scale_ue8m0:
raise ValueError("int32-packed scale buffers require scale_ue8m0=True")
aligned = num_groups % 4 == 0
return row_major, aligned
if output_s.dtype == torch.float32:
if scale_ue8m0:
raise ValueError("scale_ue8m0=True requires an int32-packed output_s")
return row_major, True
raise ValueError(f"Unsupported output_s dtype {output_s.dtype}")
@register_custom_op(
op_name="per_token_group_quant",
mutates_args=["output_q", "output_s"],
)
def _per_token_group_quant_custom_op(
input: torch.Tensor,
output_q: torch.Tensor,
output_s: torch.Tensor,
group_size: int,
scale_ue8m0: bool = False,
fuse_silu_and_mul: bool = False,
masked_m: Optional[torch.Tensor] = None,
expected_m: Optional[int] = None,
) -> None:
num_groups = output_q.shape[-1] // group_size
row_major, aligned = _infer_scale_layout(output_s, scale_ue8m0, num_groups)
module = _jit_module(
input.dtype,
output_q.dtype,
int(group_size),
bool(scale_ue8m0),
row_major,
aligned,
bool(fuse_silu_and_mul),
masked_m is not None,
is_arch_support_pdl(),
)
if masked_m is not None:
module.per_token_group_quant(
input, output_q, output_s, masked_m, int(expected_m or -1)
)
else:
module.per_token_group_quant(input, output_q, output_s)
def _allocate_outputs(
input: torch.Tensor,
group_size: int,
out_dtype: torch.dtype,
scale_ue8m0: bool,
column_major_scales: bool,
fuse_silu_and_mul: bool,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Allocate ``(output_q, output_s)`` in the requested major mode / scale
format, selected by ``(column_major_scales, scale_ue8m0)``."""
hidden = input.shape[-1] // (2 if fuse_silu_and_mul else 1)
out_shape = (*input.shape[:-1], hidden)
output_q = torch.empty(out_shape, device=input.device, dtype=out_dtype)
num_groups = hidden // group_size
if scale_ue8m0 and not column_major_scales:
# Row-major packed UE8M0: int32 [..., ceil(ng/4)] contiguous (an
# unaligned ng leaves a partially-used last int32 that the kernel zero-
# pads). The shared create_*_output_scale helper does not produce this
# layout.
output_s = torch.empty(
(*out_shape[:-1], (num_groups + 3) // 4),
device=input.device,
dtype=torch.int32,
)
else:
from sglang.kernels.ops.quantization.fp8_kernel import (
create_per_token_group_quant_fp8_output_scale,
)
output_s = create_per_token_group_quant_fp8_output_scale(
x_shape=out_shape,
device=input.device,
group_size=group_size,
column_major_scales=column_major_scales,
scale_tma_aligned=column_major_scales,
scale_ue8m0=scale_ue8m0,
)
return output_q, output_s
@debug_kernel_api
def per_token_group_quant(
input: torch.Tensor,
output_q: Optional[torch.Tensor] = None,
output_s: Optional[torch.Tensor] = None,
group_size: int = 128,
scale_ue8m0: bool = False,
fuse_silu_and_mul: bool = False,
masked_m: Optional[torch.Tensor] = None,
expected_m: Optional[int] = None,
*,
out_dtype: Optional[torch.dtype] = None,
column_major_scales: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Per-token-group quantization. Returns ``(output_q, output_s)``.
``output_q`` / ``output_s`` are optional: pass them to quantize into
caller-owned buffers, or omit both to have them allocated per ``out_dtype``
(default fp8_e4m3), ``scale_ue8m0`` and ``column_major_scales``. Either way
the two tensors are returned.
Input / output shapes:
vanilla: input [T, hidden], output_q [T, hidden]
fuse_silu_and_mul: input [T, hidden*2], output_q [T, hidden]
masked (+ above): input [E, T_pad, ...], output_q [E, T_pad, hidden],
masked_m [E] int32
``output_s`` scale layouts (inferred from a supplied buffer's dtype/strides,
or allocated to match when omitted):
float32 contiguous -> row-major fp32 scales
float32 transposed -> col-major fp32 scales (TMA-aligned view)
int32 transposed -> col-major UE8M0 bytes packed 4-per-int32
int32 contiguous -> row-major UE8M0 bytes packed 4-per-int32
The packed layouts require ``scale_ue8m0=True``.
``expected_m`` (masked only) is an optional expected-tokens-per-expert hint.
Inputs are bf16/fp16; group size is one of 16/32/64/128/256; the quant range
follows ``output_q.dtype`` (fp8_e4m3: +-448, int8: [-128, 127]).
"""
if output_q is None:
assert output_s is None
output_q, output_s = _allocate_outputs(
input,
group_size,
out_dtype or torch.float8_e4m3fn,
scale_ue8m0,
column_major_scales,
fuse_silu_and_mul,
)
else:
assert output_s is not None
assert out_dtype is None or out_dtype == output_q.dtype
_per_token_group_quant_custom_op(
input=input,
output_q=output_q,
output_s=output_s,
group_size=group_size,
scale_ue8m0=scale_ue8m0,
fuse_silu_and_mul=fuse_silu_and_mul,
masked_m=masked_m,
expected_m=expected_m,
)
return output_q, output_s
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py
+3 -134
View File
@@ -1,136 +1,5 @@
"""DEPRECATED: superseded by ``sglang.jit_kernel.per_token_group_quant`` (the
default CUDA path). No sglang runtime code may call this kernel; it is kept
only as the perf baseline for the per_token_group_quant benchmarks and its own
bit-parity tests, and will be deleted once those move to torch references.
"""
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2."""
from __future__ import annotations
from sglang.kernels.ops.quantization import _jit_per_token_group_quant_8bit_v2 as _impl
from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_module(in_dtype: torch.dtype, out_dtype: torch.dtype, use_pdl: bool) -> Module:
args = make_cpp_args(in_dtype, out_dtype, use_pdl)
return load_jit(
"per_token_group_quant_8bit_v2",
*args,
cuda_files=["gemm/per_token_group_quant_8bit_v2.cuh"],
cuda_wrappers=[
(
"per_token_group_quant_8bit_v2",
f"PerTokenGroupQuant8bitV2Kernel<{args}>::run",
)
],
# Match the AOT sgl-kernel build (-use_fast_math) so the FP8 scale
# division/rounding is bit-identical to sgl_per_token_group_quant_8bit_v2.
extra_cuda_cflags=["--use_fast_math"],
)
@register_custom_op(
op_name="per_token_group_quant_8bit_v2",
mutates_args=["output_q", "output_s"],
)
def _per_token_group_quant_8bit_v2_custom_op(
input: torch.Tensor,
output_q: torch.Tensor,
output_s: torch.Tensor,
group_size: int,
eps: float,
min_8bit: float,
max_8bit: float,
scale_ue8m0: bool = False,
fuse_silu_and_mul: bool = False,
masked_m: Optional[torch.Tensor] = None,
) -> None:
"""Opaque custom-op boundary around the JIT v2 kernel.
Registering this as a custom op (instead of calling the tvm-ffi module
directly) keeps torch.compile / piecewise-CUDA-graph from tracing into the
tvm-ffi ``Function.__call__`` (which Dynamo cannot trace). All shape-derived
scalars are computed here and passed to the kernel.
Layouts (matching the AOT v2):
vanilla: input (num_tokens, hidden), output_q (num_tokens, hidden)
fuse_silu_and_mul: input (num_tokens, hidden*2), output_q (num_tokens, hidden)
fuse_silu_and_mul+masked: input (num_experts, tokens_pad, hidden*2),
output_q (num_experts, tokens_pad, hidden), masked_m (num_experts,)
"""
masked_layout = masked_m is not None
numel = input.numel()
num_groups = numel // group_size // (2 if fuse_silu_and_mul else 1)
if num_groups == 0: # empty input -> grid 0 -> cudaErrorInvalidConfiguration
return
num_local_experts = input.shape[0] if masked_layout else 1
last = output_q.dim() - 1
is_column_major = output_s.stride(last - 1) < output_s.stride(last)
hidden_dim_num_groups = output_q.shape[last] // group_size
num_tokens_per_expert = output_q.shape[last - 1]
scale_expert_stride = output_s.stride(0) if masked_layout else 0
scale_hidden_stride = output_s.stride(last)
module = _jit_module(input.dtype, output_q.dtype, is_arch_support_pdl())
module.per_token_group_quant_8bit_v2(
input,
output_q,
output_s,
masked_m if masked_layout else input, # unused (nullptr) when not masked
int(group_size),
bool(scale_ue8m0),
bool(fuse_silu_and_mul),
bool(masked_layout),
int(num_groups),
int(num_local_experts),
bool(is_column_major),
int(hidden_dim_num_groups),
int(num_tokens_per_expert),
int(scale_expert_stride),
int(scale_hidden_stride),
)
@debug_kernel_api
def per_token_group_quant_8bit_v2(
input: torch.Tensor,
output_q: torch.Tensor,
output_s: torch.Tensor,
group_size: int,
eps: float,
min_8bit: float,
max_8bit: float,
scale_ue8m0: bool = False,
fuse_silu_and_mul: bool = False,
masked_m: Optional[torch.Tensor] = None,
) -> None:
"""JIT port of sgl_per_token_group_quant_8bit_v2 (full feature parity).
Wraps the registered custom op so torch.compile / piecewise CUDA graph treat
the tvm-ffi kernel call as an opaque boundary.
"""
_per_token_group_quant_8bit_v2_custom_op(
input=input,
output_q=output_q,
output_s=output_s,
group_size=group_size,
eps=eps,
min_8bit=min_8bit,
max_8bit=max_8bit,
scale_ue8m0=scale_ue8m0,
fuse_silu_and_mul=fuse_silu_and_mul,
masked_m=masked_m,
)
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/resolve_future_token_ids.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/rmsnorm_hf.py
+1 -1
View File
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/rope.py
+1 -1
View File
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
python/sglang/jit_kernel/set_mla_kv_buffer.py
+3 -119
View File
@@ -1,121 +1,5 @@
"""JIT TMA bulk-store path for ``set_mla_kv_buffer``.
"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer."""
Each warp scatter-writes one item's (nope, rope) row via a single
``cp.async.bulk.global.shared::cta`` store. Requires SM90+ (Hopper or later)
for the TMA bulk-store hardware. The host-side wrapper in
``sglang.srt.mem_cache.utils`` falls back to a Triton kernel for older arches.
"""
from sglang.kernels.ops.kvcache import _jit_set_mla_kv_buffer as _impl
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
if TYPE_CHECKING:
from tvm_ffi.module import Module
logger = logging.getLogger(__name__)
@cache_once
def _jit_set_mla_kv_buffer_module(
nope_bytes: int, rope_bytes: int, use_pdl: bool
) -> Module:
args = make_cpp_args(nope_bytes, rope_bytes, use_pdl)
return load_jit(
f"set_mla_kv_buffer_{nope_bytes}_{rope_bytes}",
*args,
cuda_files=["elementwise/set_mla_kv_buffer.cuh"],
cuda_wrappers=[
("set_mla_kv_buffer", f"SetMlaKVBufferKernel<{args}>::run"),
],
)
@cache_once
def can_use_set_mla_kv_buffer(nope_bytes: int, rope_bytes: int) -> bool:
"""Whether the TMA path can be used for these row byte widths.
TMA bulk store requires ``(nope_bytes + rope_bytes)`` to be a multiple of
16; both halves individually must also be a multiple of 4 (the warp-coop
smem load lower bound).
"""
if nope_bytes % 4 != 0 or rope_bytes % 4 != 0:
logger.warning(
"Unsupported nope_bytes=%d rope_bytes=%d for JIT set_mla_kv_buffer:"
" both must be multiples of 4",
nope_bytes,
rope_bytes,
)
return False
if (nope_bytes + rope_bytes) % 16 != 0:
logger.warning(
"Unsupported nope_bytes=%d rope_bytes=%d for JIT set_mla_kv_buffer:"
" (nope_bytes + rope_bytes) must be a multiple of 16 for TMA bulk store",
nope_bytes,
rope_bytes,
)
return False
try:
_jit_set_mla_kv_buffer_module(nope_bytes, rope_bytes, is_arch_support_pdl())
return True
except Exception as e: # pragma: no cover - compile-time only
logger.warning(
"Failed to load JIT set_mla_kv_buffer kernel "
"with nope_bytes=%d rope_bytes=%d: %s",
nope_bytes,
rope_bytes,
e,
)
return False
def _pick_num_warps(n_loc: int) -> int:
# Tuned on GB300: nw=4 wins below 1024 (more CTAs spread across SMs);
# nw=8 wins above (each CTA amortises the bulk-group commit better).
return 4 if n_loc <= 768 else 8
def set_mla_kv_buffer(
kv_buffer: torch.Tensor,
loc: torch.Tensor,
cache_k_nope: torch.Tensor,
cache_k_rope: torch.Tensor,
num_warps: int = 0,
) -> None:
"""Write packed [k_nope | k_rope] rows into ``kv_buffer`` at ``loc`` indices
via a TMA bulk-store. SM90+ only — the caller is expected to gate.
Shapes (last dim is treated as the row payload; any leading singleton dims
on the source tensors are flattened away):
kv_buffer: [num_pages, total_dim] or [num_pages, 1, total_dim]
cache_k_nope: [n_loc, nope_dim] or [n_loc, 1, nope_dim]
cache_k_rope: [n_loc, rope_dim] or [n_loc, 1, rope_dim]
loc: [n_loc]
"""
n_loc = loc.shape[0]
if n_loc == 0:
return
src_nope = cache_k_nope.view(n_loc, -1) if cache_k_nope.dim() != 2 else cache_k_nope
src_rope = cache_k_rope.view(n_loc, -1) if cache_k_rope.dim() != 2 else cache_k_rope
buf = kv_buffer.view(kv_buffer.shape[0], -1) if kv_buffer.dim() != 2 else kv_buffer
nope_bytes = src_nope.shape[-1] * src_nope.element_size()
rope_bytes = src_rope.shape[-1] * src_rope.element_size()
if num_warps <= 0:
num_warps = _pick_num_warps(n_loc)
module = _jit_set_mla_kv_buffer_module(
nope_bytes, rope_bytes, is_arch_support_pdl()
)
module.set_mla_kv_buffer(buf, loc, src_nope, src_rope, num_warps)
globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})
python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py
+1 -1
View File
@@ -10,8 +10,8 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit, override_jit_cuda_arch
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/jit_kernel/timestep_embedding.py
+1 -1
View File
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/transfer_mamba.py
+2 -2
View File
@@ -4,7 +4,7 @@ Provides ``transfer_kv_mamba_pf_lf`` (load: page_first -> layer_first)
and ``transfer_kv_mamba_lf_pf`` (backup: layer_first -> page_first).
Uses the shared ``load_jit`` + ``cache_once`` infrastructure from
``sglang.jit_kernel.utils`` — the same mechanism used by ``hicache.py``
``sglang.kernels.jit.utils`` — the same mechanism used by ``hicache.py``
for MHA/MLA staged write-back kernels. This ensures consistent
content-addressed caching, CUDA arch detection, and multi-worker
JIT compilation behavior across all JIT kernels.
@@ -15,8 +15,8 @@ from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
import torch
python/sglang/jit_kernel/trtllm_lora_temp/kimi_k2_moe_fused_gate.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Tuple
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/trtllm_lora_temp/moe_lora_merged_align.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
python/sglang/jit_kernel/trtllm_lora_temp/topk_softmax_pack.py
+1 -1
View File
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernels.jit.utils import cache_once, load_jit
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
python/sglang/kernels/jit/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""Internal JIT home under ``sglang.kernels`` (RFC #29630).
Mirrors the legacy ``sglang.jit_kernel`` tree; shared build/runtime
infrastructure lives in :mod:`sglang.kernels.jit.utils`. csrc / include /
operators migrate here in later phases.
"""
python/sglang/jit_kernel/utils/__init__.py → python/sglang/kernels/jit/utils/__init__.py
+4 -4
View File
@@ -1,11 +1,11 @@
"""Public interface of sglang.jit_kernel.utils."""
"""Public interface of sglang.kernels.jit.utils."""
from sglang.jit_kernel.utils.arch import (
from sglang.kernels.jit.utils.arch import (
get_jit_cuda_arch,
is_arch_support_pdl,
override_jit_cuda_arch,
)
from sglang.jit_kernel.utils.common import (
from sglang.kernels.jit.utils.common import (
cache_once,
empty_sentinel,
get_ci_test_range,
@@ -14,7 +14,7 @@ from sglang.jit_kernel.utils.common import (
lazy_register_class,
should_run_full_tests,
)
from sglang.jit_kernel.utils.compile import KERNEL_PATH, load_jit, make_cpp_args
from sglang.kernels.jit.utils.compile import KERNEL_PATH, load_jit, make_cpp_args
__all__ = [
"empty_sentinel",
python/sglang/jit_kernel/utils/arch.py → python/sglang/kernels/jit/utils/arch.py
+1 -1
View File
@@ -9,7 +9,7 @@ from typing import List
import torch
from sglang.jit_kernel.utils.common import (
from sglang.kernels.jit.utils.common import (
cache_once,
is_hip_runtime,
is_musa_runtime,
python/sglang/jit_kernel/utils/common.py → python/sglang/kernels/jit/utils/common.py
View File
python/sglang/jit_kernel/utils/compile.py → python/sglang/kernels/jit/utils/compile.py
+3 -3
View File
@@ -13,9 +13,9 @@ from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union
import torch
from sglang.jit_kernel.utils.arch import get_default_target_flags, get_jit_cuda_arch
from sglang.jit_kernel.utils.common import cache_once, is_hip_runtime
from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES
from sglang.kernels.jit.utils.arch import get_default_target_flags, get_jit_cuda_arch
from sglang.kernels.jit.utils.common import cache_once, is_hip_runtime
from sglang.kernels.jit.utils.deps import REGISTERED_DEPENDENCIES
if TYPE_CHECKING:
from tvm_ffi import Module
python/sglang/jit_kernel/utils/deps.py → python/sglang/kernels/jit/utils/deps.py
View File
python/sglang/kernels/ops/activation/__init__.py
+4 -4
View File
@@ -31,7 +31,7 @@ _HIP = frozenset({CapabilityRequirement.HIP})
# — the canonical OR-semantics case that a device-baked backend name couldn't.
_CUDA_HIP = frozenset({CapabilityRequirement.CUDA, CapabilityRequirement.HIP})
# JIT before AOT to match the production path (srt/layers/activation.py imports
# from sglang.jit_kernel.activation on CUDA); auto-selection must not invert it.
# from sglang.kernels.ops.activation._jit_activation on CUDA); auto-selection must not invert it.
_ACT_PRIORITY = (
KernelBackend.JIT,
KernelBackend.AOT,
@@ -82,7 +82,7 @@ class _GatedActivationOp(BaseFusedOp):
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
import sglang.jit_kernel.activation as jit_activation
import sglang.kernels.ops.activation._jit_activation as jit_activation
return getattr(jit_activation, self.kernel_attr)(
input, out, expert_ids, expert_step
@@ -183,7 +183,7 @@ class GeluTanhAndMulOp(_GatedActivationOp):
class ReLU2Op(BaseFusedOp):
"""``out = relu(input) ** 2`` (single-input, not gated).
The real kernel is the CUDA JIT path (``sglang.jit_kernel.activation.relu2``,
The real kernel is the CUDA JIT path (``sglang.kernels.ops.activation._jit_activation.relu2``,
used in production on CUDA); elsewhere the torch reference runs.
"""
@@ -214,7 +214,7 @@ class ReLU2Op(BaseFusedOp):
def forward_jit(
self, input: torch.Tensor, out: Optional[torch.Tensor] = None
) -> torch.Tensor:
from sglang.jit_kernel.activation import relu2
from sglang.kernels.ops.activation._jit_activation import relu2
result = relu2(input)
if out is None:
python/sglang/kernels/ops/activation/_jit_activation.py
+168
View File
@@ -0,0 +1,168 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernels.jit.utils import (
cache_once,
get_jit_cuda_arch,
is_arch_support_pdl,
is_hip_runtime,
load_jit,
make_cpp_args,
)
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
def _fast_math_flags() -> list[str]:
# Mirrors sgl-kernel's CMake policy: fast-math on SM90, precise on
# SM100+ (Blackwell needs bit-exact expf), off on HIP (clang rejects).
if is_hip_runtime():
return []
if get_jit_cuda_arch().major >= 10:
return []
return ["--use_fast_math"]
@cache_once
def _jit_activation_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype, is_arch_support_pdl())
return load_jit(
"activation",
*args,
cuda_files=["elementwise/activation.cuh"],
extra_cuda_cflags=_fast_math_flags(),
cuda_wrappers=[
("run_activation", f"ActivationKernel<{args}>::run_activation"),
(
"run_activation_filtered",
f"ActivationKernel<{args}>::run_activation_filtered",
),
(
"run_unary_activation",
f"ActivationKernel<{args}>::run_unary_activation",
),
],
)
SUPPORTED_ACTIVATIONS = {"silu", "gelu", "gelu_tanh"}
SUPPORTED_UNARY_ACTIVATIONS = {"relu2"}
@register_custom_op(mutates_args=["out"])
def _run_activation_inplace(
op_name: str, input: torch.Tensor, out: torch.Tensor
) -> None:
hidden_size = input.shape[-1] // 2
module = _jit_activation_module(input.dtype)
input_2d = input.view(-1, hidden_size * 2)
out_2d = out.view(-1, hidden_size)
module.run_activation(input_2d, out_2d, op_name)
@register_custom_op(mutates_args=["out"])
def _run_activation_filtered_inplace(
op_name: str,
input: torch.Tensor,
out: torch.Tensor,
expert_ids: torch.Tensor,
expert_step: int,
) -> None:
hidden_size = input.shape[-1] // 2
module = _jit_activation_module(input.dtype)
input_2d = input.view(-1, hidden_size * 2)
out_2d = out.view(-1, hidden_size)
module.run_activation_filtered(input_2d, out_2d, expert_ids, expert_step, op_name)
def run_activation(
op_name: str,
input: torch.Tensor,
out: Optional[torch.Tensor],
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
"""Apply ``op_name`` activation followed by element-wise multiplication.
When ``expert_ids`` is provided, output rows are skipped for tokens whose
routed expert id is ``-1``. ``expert_step`` is 1 for per-token routing and
``BLOCK_SIZE_M`` for sorted/TMA routing — i.e. ``expert_ids[token_id //
expert_step]`` is consulted before computing each row.
"""
assert op_name in SUPPORTED_ACTIVATIONS, f"Unsupported activation: {op_name}"
hidden_size = input.shape[-1] // 2
if out is None:
out = input.new_empty(*input.shape[:-1], hidden_size)
if expert_ids is None:
_run_activation_inplace(op_name, input, out)
else:
_run_activation_filtered_inplace(op_name, input, out, expert_ids, expert_step)
return out
@register_custom_op(mutates_args=["out"])
def _run_unary_activation_inplace(
op_name: str, input: torch.Tensor, out: torch.Tensor
) -> None:
last = input.shape[-1]
module = _jit_activation_module(input.dtype)
module.run_unary_activation(input.view(-1, last), out.view(-1, last), op_name)
def run_unary_activation(
op_name: str,
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Apply a standalone (non-gated) element-wise activation: ``out = act(input)``.
Unlike :func:`run_activation`, there is no gate/up split — ``input`` and
``out`` share the same shape.
"""
assert (
op_name in SUPPORTED_UNARY_ACTIVATIONS
), f"Unsupported unary activation: {op_name}"
if out is None:
out = torch.empty_like(input)
_run_unary_activation_inplace(op_name, input, out)
return out
def relu2(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Squared ReLU: ``out = max(0, input) ** 2`` (element-wise)."""
return run_unary_activation("relu2", input, out)
def silu_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("silu", input, out, expert_ids, expert_step)
def gelu_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("gelu", input, out, expert_ids, expert_step)
def gelu_tanh_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("gelu_tanh", input, out, expert_ids, expert_step)
python/sglang/kernels/ops/attention/dsv4/compress_c128_hip.py
+1 -1
View File
@@ -14,7 +14,7 @@ from sglang.jit_kernel.dsv4 import (
CompressorDecodePlan,
CompressorPrefillPlan,
)
from sglang.jit_kernel.utils import is_hip_runtime
from sglang.kernels.jit.utils import is_hip_runtime
_is_hip = is_hip_runtime()
python/sglang/kernels/ops/attention/fla/layernorm_gated.py
+1 -1
View File
@@ -15,7 +15,7 @@ import triton
import triton.language as tl
from einops import rearrange
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
python/sglang/kernels/ops/attention/utils.py
+1 -1
View File
@@ -2,7 +2,7 @@ import torch
import triton
import triton.language as tl
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.kernels.ops.attention.pad import (
pad_sequence_with_mask as pad_sequence_with_mask,
)
python/sglang/kernels/ops/gemm/__init__.py
+2 -2
View File
@@ -59,7 +59,7 @@ register_kernel(
KernelSpec(
op="gemm.dsv3_fused_a_gemm",
backend=KernelBackend.JIT,
target="sglang.jit_kernel.dsv3_fused_a_gemm:dsv3_fused_a_gemm",
target="sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm:dsv3_fused_a_gemm",
capabilities=_CUDA,
format_signature=FormatSignature(
supported_dtypes=("bfloat16",),
@@ -72,7 +72,7 @@ register_kernel(
KernelSpec(
op="gemm.dsv3_router_gemm",
backend=KernelBackend.JIT,
target="sglang.jit_kernel.dsv3_router_gemm:dsv3_router_gemm",
target="sglang.kernels.ops.gemm._jit_dsv3_router_gemm:dsv3_router_gemm",
capabilities=_CUDA,
format_signature=FormatSignature(
supported_dtypes=("bfloat16",),
python/sglang/kernels/ops/gemm/_jit_dsv3_fused_a_gemm.py
+90
View File
@@ -0,0 +1,90 @@
"""
JIT kernel for DeepSeek V3 fused QKV-A GEMM (min-latency).
Runtime-compiled CUDA C++ kernel for SM90+ (Hopper) GPUs.
Shapes: hd_in a multiple of 256, hd_out a multiple of 16, num_tokens 1-16, bfloat16.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.srt.utils.common import direct_register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_dsv3_fused_a_gemm_module(hd_in: int, hd_out: int, use_pdl: bool) -> Module:
args = make_cpp_args(hd_in, hd_out, use_pdl)
return load_jit(
"dsv3_fused_a_gemm",
*args,
cuda_files=["gemm/dsv3_fused_a_gemm.cuh"],
cuda_wrappers=[
("dsv3_fused_a_gemm", f"DSV3FusedAGemmKernel<{args}>::run"),
],
)
def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
assert mat_a.stride(1) == 1, "mat_a must be row-major [M, K]"
output = torch.empty(
(mat_a.shape[0], mat_b.shape[1]),
device=mat_a.device,
dtype=mat_a.dtype,
)
module = _jit_dsv3_fused_a_gemm_module(
mat_a.shape[1], mat_b.shape[1], is_arch_support_pdl()
)
module.dsv3_fused_a_gemm(mat_a, mat_b, output)
return output
def _dsv3_fused_a_gemm_fake(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
return mat_a.new_empty((mat_a.shape[0], mat_b.shape[1]), dtype=torch.bfloat16)
direct_register_custom_op(
op_name="jit_dsv3_fused_a_gemm",
op_func=_dsv3_fused_a_gemm_run,
mutates_args=[],
fake_impl=_dsv3_fused_a_gemm_fake,
)
@debug_kernel_api
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 kernel (JIT variant).
Args:
mat_a: Input tensor of shape [num_tokens, hd_in], bfloat16, row-major.
hd_in must be a multiple of 256 and num_tokens in [1, 16].
mat_b: Weight tensor of shape [hd_in, hd_out], bfloat16, column-major
(i.e. ``weight.T`` of a row-major [hd_out, hd_in] weight).
hd_out must be a multiple of 16.
output: Optional pre-allocated output tensor of shape [num_tokens, hd_out].
Returns:
Output tensor of shape [num_tokens, hd_out].
"""
result = torch.ops.sglang.jit_dsv3_fused_a_gemm(mat_a, mat_b)
if output is not None:
output.copy_(result)
return output
return result
python/sglang/kernels/ops/gemm/_jit_dsv3_router_gemm.py
+92
View File
@@ -0,0 +1,92 @@
"""
JIT kernel for DeepSeek V3 router GEMM.
Runtime-compiled CUDA C++ kernel for SM90+ (Hopper) GPUs.
Supports num_experts in {256, 384}, hidden_dim a multiple of 1024, num_tokens 1-16.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_dsv3_router_gemm_module(
num_experts: int,
hidden_dim: int,
use_pdl: bool,
out_float: bool,
) -> Module:
args = make_cpp_args(num_experts, hidden_dim, use_pdl, out_float)
return load_jit(
"dsv3_router_gemm",
*args,
cuda_files=["gemm/dsv3_router_gemm.cuh"],
cuda_wrappers=[
("dsv3_router_gemm", f"DSV3RouterGemmKernel<{args}>::run"),
],
)
@register_custom_op(
op_name="dsv3_router_gemm",
mutates_args=["output"],
)
def _dsv3_router_gemm_custom_op(
hidden_states: torch.Tensor,
router_weights: torch.Tensor,
output: torch.Tensor,
) -> None:
num_experts = router_weights.shape[0]
hidden_dim = hidden_states.shape[1]
out_float = output.dtype == torch.float32
module = _jit_dsv3_router_gemm_module(
num_experts, hidden_dim, is_arch_support_pdl(), out_float
)
module.dsv3_router_gemm(hidden_states, router_weights, output)
return None
@debug_kernel_api
def dsv3_router_gemm(
hidden_states: torch.Tensor,
router_weights: torch.Tensor,
out_dtype: torch.dtype = torch.bfloat16,
output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
DeepSeek V3 router GEMM kernel (JIT variant).
Args:
hidden_states: Input tensor of shape [num_tokens, hidden_dim], bfloat16.
hidden_dim must be a multiple of 1024 and num_tokens in [1, 16].
router_weights: Weight tensor of shape [num_experts, hidden_dim], bfloat16.
out_dtype: Output dtype, either torch.bfloat16 or torch.float32.
output: Optional pre-allocated output tensor.
Returns:
Output tensor of shape [num_tokens, num_experts].
"""
if output is None:
output = torch.empty(
hidden_states.shape[0],
router_weights.shape[0],
device=hidden_states.device,
dtype=out_dtype,
)
_dsv3_router_gemm_custom_op(hidden_states, router_weights, output)
return output
python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py
+1 -1
View File
@@ -1,7 +1,7 @@
import triton
import triton.language as tl
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.kernels.jit.utils import is_arch_support_pdl
def get_pdl_launch_metadata() -> tuple[bool, dict]:
python/sglang/kernels/ops/kvcache/_jit_set_mla_kv_buffer.py
+121
View File
@@ -0,0 +1,121 @@
"""JIT TMA bulk-store path for ``set_mla_kv_buffer``.
Each warp scatter-writes one item's (nope, rope) row via a single
``cp.async.bulk.global.shared::cta`` store. Requires SM90+ (Hopper or later)
for the TMA bulk-store hardware. The host-side wrapper in
``sglang.srt.mem_cache.utils`` falls back to a Triton kernel for older arches.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
if TYPE_CHECKING:
from tvm_ffi.module import Module
logger = logging.getLogger(__name__)
@cache_once
def _jit_set_mla_kv_buffer_module(
nope_bytes: int, rope_bytes: int, use_pdl: bool
) -> Module:
args = make_cpp_args(nope_bytes, rope_bytes, use_pdl)
return load_jit(
f"set_mla_kv_buffer_{nope_bytes}_{rope_bytes}",
*args,
cuda_files=["elementwise/set_mla_kv_buffer.cuh"],
cuda_wrappers=[
("set_mla_kv_buffer", f"SetMlaKVBufferKernel<{args}>::run"),
],
)
@cache_once
def can_use_set_mla_kv_buffer(nope_bytes: int, rope_bytes: int) -> bool:
"""Whether the TMA path can be used for these row byte widths.
TMA bulk store requires ``(nope_bytes + rope_bytes)`` to be a multiple of
16; both halves individually must also be a multiple of 4 (the warp-coop
smem load lower bound).
"""
if nope_bytes % 4 != 0 or rope_bytes % 4 != 0:
logger.warning(
"Unsupported nope_bytes=%d rope_bytes=%d for JIT set_mla_kv_buffer:"
" both must be multiples of 4",
nope_bytes,
rope_bytes,
)
return False
if (nope_bytes + rope_bytes) % 16 != 0:
logger.warning(
"Unsupported nope_bytes=%d rope_bytes=%d for JIT set_mla_kv_buffer:"
" (nope_bytes + rope_bytes) must be a multiple of 16 for TMA bulk store",
nope_bytes,
rope_bytes,
)
return False
try:
_jit_set_mla_kv_buffer_module(nope_bytes, rope_bytes, is_arch_support_pdl())
return True
except Exception as e: # pragma: no cover - compile-time only
logger.warning(
"Failed to load JIT set_mla_kv_buffer kernel "
"with nope_bytes=%d rope_bytes=%d: %s",
nope_bytes,
rope_bytes,
e,
)
return False
def _pick_num_warps(n_loc: int) -> int:
# Tuned on GB300: nw=4 wins below 1024 (more CTAs spread across SMs);
# nw=8 wins above (each CTA amortises the bulk-group commit better).
return 4 if n_loc <= 768 else 8
def set_mla_kv_buffer(
kv_buffer: torch.Tensor,
loc: torch.Tensor,
cache_k_nope: torch.Tensor,
cache_k_rope: torch.Tensor,
num_warps: int = 0,
) -> None:
"""Write packed [k_nope | k_rope] rows into ``kv_buffer`` at ``loc`` indices
via a TMA bulk-store. SM90+ only — the caller is expected to gate.
Shapes (last dim is treated as the row payload; any leading singleton dims
on the source tensors are flattened away):
kv_buffer: [num_pages, total_dim] or [num_pages, 1, total_dim]
cache_k_nope: [n_loc, nope_dim] or [n_loc, 1, nope_dim]
cache_k_rope: [n_loc, rope_dim] or [n_loc, 1, rope_dim]
loc: [n_loc]
"""
n_loc = loc.shape[0]
if n_loc == 0:
return
src_nope = cache_k_nope.view(n_loc, -1) if cache_k_nope.dim() != 2 else cache_k_nope
src_rope = cache_k_rope.view(n_loc, -1) if cache_k_rope.dim() != 2 else cache_k_rope
buf = kv_buffer.view(kv_buffer.shape[0], -1) if kv_buffer.dim() != 2 else kv_buffer
nope_bytes = src_nope.shape[-1] * src_nope.element_size()
rope_bytes = src_rope.shape[-1] * src_rope.element_size()
if num_warps <= 0:
num_warps = _pick_num_warps(n_loc)
module = _jit_set_mla_kv_buffer_module(
nope_bytes, rope_bytes, is_arch_support_pdl()
)
module.set_mla_kv_buffer(buf, loc, src_nope, src_rope, num_warps)
python/sglang/kernels/ops/kvcache/mla_buffer.py
+3 -3
View File
@@ -4,7 +4,7 @@ import torch
import triton
import triton.language as tl
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.srt.runtime_context import get_parallel
@@ -116,10 +116,10 @@ def set_mla_kv_buffer_triton(
Name retained for caller compatibility; the implementation is no longer
Triton-only.
"""
from sglang.jit_kernel.set_mla_kv_buffer import (
from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import (
can_use_set_mla_kv_buffer,
)
from sglang.jit_kernel.set_mla_kv_buffer import (
from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import (
set_mla_kv_buffer as jit_set_mla_kv_buffer,
)
python/sglang/kernels/ops/layernorm/__init__.py
+4 -2
View File
@@ -113,7 +113,7 @@ class RMSNormOp(BaseFusedOp):
) -> torch.Tensor:
import torch
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm
if out is None:
out = torch.empty_like(input)
@@ -227,7 +227,9 @@ class FusedAddRMSNormOp(BaseFusedOp):
eps: float = 1e-6,
enable_pdl: Optional[bool] = None,
) -> None:
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.kernels.ops.layernorm._jit_norm import (
fused_add_rmsnorm as jit_fused_add_rmsnorm,
)
return jit_fused_add_rmsnorm(input, residual, weight, eps)

Some files were not shown because too many files have changed in this diff Show More