diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index 739651de2..343039c57 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -186,9 +186,7 @@ LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params); ## Step 0 (optional): Generate a `.clangd` config for better IDE support ```bash -python -m sglang.jit_kernel -h # for verbose help info about clangd configuration python -m sglang.jit_kernel -python -m sglang.jit_kernel --dep cutlass flashinfer # with cutlass/flashinfer dependency ``` --- @@ -200,8 +198,6 @@ Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`. The implementation fully uses the project abstractions described above: ```cpp -// NOTE: Comments for headers are not common in practice. -// It is only shown here for tutorial purposes to highlight the key abstractions. #include // For TensorMatcher, SymbolicSize, SymbolicDevice #include // For dtype_trait, fp16_t, bf16_t, fp32_t #include // For RuntimeCheck, div_ceil @@ -219,7 +215,7 @@ namespace { // kVecN = number of elements per vector load (e.g. 8 for fp16) // factor = runtime scale factor // ---------------------------------------------------------------- -template +template __global__ void scale_kernel(T* __restrict__ dst, const T* __restrict__ src, float factor, @@ -227,10 +223,6 @@ __global__ void scale_kernel(T* __restrict__ dst, using vec_t = device::AlignedVector; const uint32_t n_vecs = n_total / kVecN; - // If using PDL, wait for primary kernel before any global memory load. - // This is NOT a synchronization point, which means some threads can early exit before this. - device::PDLWaitPrimary(); - // --- vectorised body --- const uint32_t vec_stride = blockDim.x * gridDim.x; for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x; @@ -253,16 +245,12 @@ __global__ void scale_kernel(T* __restrict__ dst, i += scalar_stride) { dst[base + i] = static_cast(static_cast(src[base + i]) * factor); } - - // If using PDL, signal for the secondary kernel to start after all threads have finished - // This is NOT a synchronization point, which means some threads can early exit before this. - device::PDLTriggerSecondary(); } // ---------------------------------------------------------------- // Launcher: validates tensors, selects vector width, launches kernel // ---------------------------------------------------------------- -template +template void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { using namespace host; @@ -283,12 +271,8 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { RuntimeCheck(n > 0, "scale: num_elements must be > 0, got ", n); // 2. Choose vector width for 128-bit loads (16 bytes) - // fp16/bf16: 8 elements x 2 bytes = 16 bytes - // fp32: 4 elements x 4 bytes = 16 bytes - // We encourage using `device::kMaxVecBytes`, which will change according to - // the target architecture and can enable 256-bit vectorization on SM100+ if desired. - // But 128-bit is more commonly adapted for better compatibility, - // so it's still ok to hardcode 16 here just for simplicity. + // fp16/bf16: 8 elements × 2 bytes = 16 bytes + // fp32: 4 elements × 4 bytes = 16 bytes constexpr int kVecN = 16 / sizeof(T); const uint32_t n_work_items = div_ceil(n, static_cast(kVecN)); @@ -296,10 +280,8 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { constexpr uint32_t kBlockSize = 256; const uint32_t grid = div_ceil(n_work_items, kBlockSize); - // PDL feature is 100% optional. Without `enable_pdl`, the code should still be correct. - // Try to enable it if profiling shows that it can benefit the performance of this kernel. - LaunchKernel(grid, kBlockSize, device).enable_pdl(kUsePDL)( - scale_kernel, + LaunchKernel(grid, kBlockSize, device)( + scale_kernel, static_cast(dst.data_ptr()), static_cast(src.data_ptr()), factor, @@ -320,8 +302,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { - Prefer passing runtime scalars like `factor` directly unless compile-time specialisation is genuinely required - `fp16_t` / `bf16_t` / `fp32_t` are the project's type aliases (from `utils.cuh`) - `device::cast` or `dtype_trait::from(val)` for cross-type conversions -- `device::math::` functions for device math instead of bare `__` intrinsics if possible. -- Try to use `PDL` feature. In some cases, this will benefit the performance. +- `device::math::` functions for device math instead of bare `__` intrinsics --- @@ -336,12 +317,7 @@ from typing import TYPE_CHECKING import torch -from sglang.jit_kernel.utils import ( - cache_once, - is_arch_support_pdl, - load_jit, - make_cpp_args, -) +from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @@ -350,7 +326,7 @@ if TYPE_CHECKING: @cache_once def _jit_scale_module(dtype: torch.dtype) -> Module: """Compile and cache the JIT scale module for a given dtype.""" - args = make_cpp_args(dtype, is_arch_support_pdl()) + args = make_cpp_args(dtype) return load_jit( "scale", *args, @@ -375,16 +351,24 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> ------- Scaled tensor (dst = src * factor). """ - # DO NOT add too much proactive validation here. - # Keep the Python wrapper thin, only enforce the preconditions - # that the current JIT/FFI path (C++ side) does not reject on its own. + if not src.is_cuda: + raise RuntimeError("src must be a CUDA tensor") if src.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise RuntimeError( f"Unsupported dtype {src.dtype}. Supported: float16, bfloat16, float32" ) if out is None: out = torch.empty_like(src) + else: + if out.shape != src.shape: + raise RuntimeError("out shape must match src") + if out.dtype != src.dtype: + raise RuntimeError("out dtype must match src") + if out.device != src.device: + raise RuntimeError("out device must match src") + # Keep the Python wrapper thin, but still enforce the basic preconditions + # that the current JIT/FFI path does not reject safely on its own. module = _jit_scale_module(src.dtype) module.scale(out, src, factor) return out @@ -397,7 +381,6 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> - Only include compile-time specialisation knobs in the build marker; runtime values like `factor` should stay runtime unless the kernel truly needs templating - `cuda_wrappers`: `(export_name, kernel_symbol)` — `export_name` is called from Python - `make_cpp_args(dtype, ...)` converts `torch.dtype` to C++ type alias: -- `is_arch_support_pdl()` checks if the current architecture supports PDL, which is typically passed as a template argument to the kernel. - Keep Python launchers thin, but still validate the basic invariants (`is_cuda`, supported dtype, `out` metadata). In the current JIT/FFI path, invalid tensors are not always rejected safely before launch | `torch.dtype` | C++ type | @@ -410,8 +393,6 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> ## Step 3 (optional): Tune JIT build flags -If your kernel uses some math functions like `expf` or `sinf`, consider enabling `--use_fast_math` for better performance (with a potential precision tradeoff): - ```python return load_jit( "scale", @@ -433,7 +414,7 @@ if torch.cuda.get_device_capability()[0] < 9: ## Step 4: Write tests (required) -JIT kernel tests live under `python/sglang/jit_kernel/tests/`. **CI does not run `pytest` in that directory directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` there (and every `bench_*.py` under `benchmark/`), collects `register_*_ci(...)` calls by **statically parsing each file's AST**, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check. +JIT kernel tests live under `python/sglang/jit_kernel/tests/`. **CI does not run `pytest` in that directory directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` there (and every `bench_*.py` under `benchmark/`), collects `register_*_ci(...)` calls by **statically parsing each file’s AST**, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check. - **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `stage-b-kernel-unit-1-gpu-large` (see `.github/workflows/pr-test-jit-kernel.yml`: `python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-large`). - **Nightly kernel suite**: `nightly-kernel-1-gpu` with `--nightly` — typically used with `SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1` in CI for expanded parameter grids (see `python/sglang/jit_kernel/utils.py` → `should_run_full_tests` / `get_ci_test_range`). Wired in `.github/workflows/nightly-test-nvidia.yml` (e.g. `python3 run_suite.py --hw cuda --suite nightly-kernel-1-gpu --nightly --continue-on-error`). diff --git a/python/sglang/jit_kernel/activation.py b/python/sglang/jit_kernel/activation.py deleted file mode 100644 index 96e96f06f..000000000 --- a/python/sglang/jit_kernel/activation.py +++ /dev/null @@ -1,66 +0,0 @@ -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.srt.utils.custom_op import register_custom_op - -if TYPE_CHECKING: - from tvm_ffi.module import Module - - -@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=["--use_fast_math"], - cuda_wrappers=[ - ("run_activation", f"ActivationKernel<{args}>::run_activation"), - ], - ) - - -SUPPORTED_ACTIVATIONS = {"silu", "gelu", "gelu_tanh"} - - -@register_custom_op(mutates_args=["out"], out_shape="input") -def run_activation( - op_name: str, input: torch.Tensor, out: Optional[torch.Tensor] -) -> torch.Tensor: - assert op_name in SUPPORTED_ACTIVATIONS, f"Unsupported activation: {op_name}" - - if out is None: - out = input.new_empty(*input.shape[:-1], input.shape[-1] // 2) - module = _jit_activation_module(input.dtype) - input_2d = input.view(-1, input.shape[-1]) - out_2d = out.view(-1, out.shape[-1]) - module.run_activation(input_2d, out_2d, op_name) - return out - - -def silu_and_mul( - input: torch.Tensor, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return run_activation("silu", input, out) - - -def gelu_and_mul( - input: torch.Tensor, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return run_activation("gelu", input, out) - - -def gelu_tanh_and_mul( - input: torch.Tensor, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return run_activation("gelu_tanh", input, out) diff --git a/python/sglang/jit_kernel/benchmark/bench_activation.py b/python/sglang/jit_kernel/benchmark/bench_activation.py deleted file mode 100644 index 2caac5551..000000000 --- a/python/sglang/jit_kernel/benchmark/bench_activation.py +++ /dev/null @@ -1,86 +0,0 @@ -import itertools - -import torch -import torch.nn.functional as F -import triton -import triton.testing -from sgl_kernel import gelu_and_mul as gelu_and_mul_aot -from sgl_kernel import gelu_tanh_and_mul as gelu_tanh_and_mul_aot -from sgl_kernel import silu_and_mul as silu_and_mul_aot - -from sglang.jit_kernel.activation import gelu_and_mul as gelu_and_mul_jit -from sglang.jit_kernel.activation import gelu_tanh_and_mul as gelu_tanh_and_mul_jit -from sglang.jit_kernel.activation import silu_and_mul as silu_and_mul_jit -from sglang.jit_kernel.benchmark.utils import ( - DEFAULT_DEVICE, - DEFAULT_DTYPE, - get_benchmark_range, - run_benchmark, -) -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=30, suite="stage-b-kernel-benchmark-1-gpu-large") - - -@torch.compile -def silu_and_mul(input: torch.Tensor) -> torch.Tensor: - lhs, rhs = input.split(input.shape[-1] // 2, dim=-1) - return F.silu(lhs) * rhs - - -@torch.compile -def gelu_and_mul(input: torch.Tensor) -> torch.Tensor: - lhs, rhs = input.split(input.shape[-1] // 2, dim=-1) - return F.gelu(lhs, approximate="none") * rhs - - -@torch.compile -def gelu_tanh_and_mul(input: torch.Tensor) -> torch.Tensor: - lhs, rhs = input.split(input.shape[-1] // 2, dim=-1) - return F.gelu(lhs, approximate="tanh") * rhs - - -OPS = { - "silu": (silu_and_mul_aot, silu_and_mul_jit, silu_and_mul), - "gelu": (gelu_and_mul_aot, gelu_and_mul_jit, gelu_and_mul), - "gelu_tanh": (gelu_tanh_and_mul_aot, gelu_tanh_and_mul_jit, gelu_tanh_and_mul), -} -BS_LIST = get_benchmark_range(full_range=[2**x for x in range(0, 15)], ci_range=[8]) -DIM_LIST = get_benchmark_range(full_range=[1024, 4096, 6144, 8192], ci_range=[4096]) -CONFIGS = list(itertools.product(OPS, DIM_LIST, BS_LIST)) -NUM_LAYERS = 4 # to eliminate L2 effect - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["op_name", "dim", "batch_size"], - x_vals=CONFIGS, - line_arg="provider", - line_vals=["aot", "jit", "torch"], - line_names=["AOT (sgl-kernel)", "JIT (jit_kernel)", "torch.compile"], - styles=[("blue", "--"), ("orange", "-"), ("green", "-")], - ylabel="us", - plot_name="activation-aot-vs-jit", - args={}, - ) -) -def benchmark(op_name: str, dim: int, batch_size: int, provider: str): - x = torch.randn( - NUM_LAYERS, - batch_size, - 2 * dim, - dtype=DEFAULT_DTYPE, - device=DEFAULT_DEVICE, - ) - aot_op, jit_op, torch_op = OPS[op_name] - fn = {"aot": aot_op, "jit": jit_op, "torch": torch_op}[provider] - - def f(): - for i in range(NUM_LAYERS): - fn(x[i]) - - return run_benchmark(f, scale=NUM_LAYERS) - - -if __name__ == "__main__": - benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/elementwise/activation.cuh b/python/sglang/jit_kernel/csrc/elementwise/activation.cuh deleted file mode 100644 index 0cd2e4b6b..000000000 --- a/python/sglang/jit_kernel/csrc/elementwise/activation.cuh +++ /dev/null @@ -1,137 +0,0 @@ -#include -#include - -#include -#include -#include -#include - -#include - -#include -#include -#include -#include - -namespace { - -enum class ActivationKind : uint32_t { - kSiLU, - kGELU, - kGELUTanh, -}; - -template -SGL_DEVICE T apply_activation(T x) { - const float x_f32 = device::cast(x); - float y_f32 = 0.0f; - - if constexpr (kAct == ActivationKind::kSiLU) { - y_f32 = x_f32 / (1.0f + expf(-x_f32)); - } else if constexpr (kAct == ActivationKind::kGELU) { - constexpr auto kSqrt1Over2 = 0.7071067811865475f; - y_f32 = x_f32 * (0.5f * (1.0f + erff(x_f32 * kSqrt1Over2))); - } else if constexpr (kAct == ActivationKind::kGELUTanh) { - constexpr auto kGeluTanhAlpha = 0.044715f; - constexpr auto kGeluTanhBeta = 0.7978845608028654f; - const float x_cube = x_f32 * x_f32 * x_f32; - const float cdf = 0.5f * (1.0f + tanhf(kGeluTanhBeta * (x_f32 + kGeluTanhAlpha * x_cube))); - y_f32 = x_f32 * cdf; - } else { - static_assert(host::dependent_false_v, "unsupported activation kind"); - } - - return device::cast(y_f32); -} - -struct ActivationParams { - const void* __restrict__ input; - void* __restrict__ out; - uint32_t hidden_dim; - uint32_t num_tokens; -}; - -template -__global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams params) { - using namespace device; - constexpr auto kVecSize = kMaxVecBytes / sizeof(T); - using vec_t = AlignedVector; - const auto num_vecs = params.hidden_dim / kVecSize; // per token - const auto tid = blockIdx.x * blockDim.x + threadIdx.x; - const auto token_id = tid / num_vecs; - - if (token_id >= params.num_tokens) return; - const auto offset = tid % num_vecs; - const auto input_offset = token_id * (num_vecs * 2) + offset; - const auto output_offset = tid; - PDLWaitPrimary(); - const auto gate = device::load_as(params.input, input_offset); - const auto up = device::load_as(params.input, input_offset + num_vecs); - vec_t out; -#pragma unroll - for (int i = 0; i < kVecSize; ++i) { - out[i] = apply_activation(gate[i]) * up[i]; - } - device::store_as(params.out, out, output_offset); - PDLTriggerSecondary(); -} - -template -struct ActivationKernel { - static constexpr auto kVecSize = device::kMaxVecBytes / sizeof(T); - static constexpr auto kBlockSize = 256u; - - template - static constexpr auto activation_kernel = act_and_mul_kernel; - - static_assert(device::kMaxVecBytes % sizeof(T) == 0, "unsupported data type"); - static void run_activation(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) { - using namespace host; - - auto N = SymbolicSize{"num_tokens"}; - auto D_in = SymbolicSize{"input_width"}; - auto D_out = SymbolicSize{"output_width"}; - auto device_ = SymbolicDevice{}; - device_.set_options(); - - TensorMatcher({N, D_out}) // - .with_dtype() - .with_device(device_) - .verify(out); - TensorMatcher({N, D_in}) // - .with_dtype() - .with_device(device_) - .verify(input); - - const auto hidden_size = D_out.unwrap(); - const auto num_tokens = static_cast(N.unwrap()); - const auto device = device_.unwrap(); - RuntimeCheck(hidden_size * 2 == D_in.unwrap(), "invalid activation dimension"); - RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible by vector size"); - const auto kernel = [&]() -> decltype(activation_kernel) { - if (type == "silu") { - return activation_kernel; - } else if (type == "gelu") { - return activation_kernel; - } else if (type == "gelu_tanh") { - return activation_kernel; - } else { - Panic("unsupported activation type: ", type); - } - return nullptr; - }(); - // only get once to avoid overhead - const auto num_total_items = num_tokens * (hidden_size / kVecSize); - RuntimeCheck(num_total_items <= std::numeric_limits::max(), "too many items for 32-bit indexing"); - const auto num_blocks = div_ceil(static_cast(num_total_items), kBlockSize); - const auto params = ActivationParams{ - .input = input.data_ptr(), - .out = out.data_ptr(), - .hidden_dim = hidden_size, - .num_tokens = num_tokens, - }; - LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params); - } -}; - -} // namespace diff --git a/python/sglang/jit_kernel/tests/test_activation.py b/python/sglang/jit_kernel/tests/test_activation.py deleted file mode 100644 index 075f3e618..000000000 --- a/python/sglang/jit_kernel/tests/test_activation.py +++ /dev/null @@ -1,79 +0,0 @@ -import sys - -import pytest -import torch -import torch.nn.functional as F - -from sglang.jit_kernel.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul -from sglang.jit_kernel.utils import get_ci_test_range -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=20, suite="stage-b-kernel-unit-1-gpu-large") -register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) - - -OPS = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} -DTYPES = [torch.float16, torch.bfloat16, torch.float32] -SHAPES = get_ci_test_range( - full_range=[ - (7, 16), - (83, 1024), - (3, 5, 16), - (2, 3, 512), - (1, 17, 4096), - ], - ci_range=[(7, 16), (2, 3, 512)], -) - - -def _reference(op_name: str, x: torch.Tensor) -> torch.Tensor: - d = x.shape[-1] // 2 - lhs = x[..., :d].float() - rhs = x[..., d:] - if op_name == "silu": - act = F.silu(lhs) - elif op_name == "gelu": - act = F.gelu(lhs, approximate="none") - else: - act = F.gelu(lhs, approximate="tanh") - return act.to(dtype=x.dtype) * rhs - - -def _tolerances(dtype: torch.dtype) -> tuple[float, float]: - if dtype == torch.float32: - return 1e-4, 1e-4 - return 1e-2, 1e-2 - - -@pytest.mark.parametrize("op_name", OPS) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("shape", SHAPES) -def test_activation_correctness( - op_name: str, dtype: torch.dtype, shape: tuple[int, ...] -) -> None: - torch.manual_seed(42) - x = torch.randn(shape, dtype=dtype, device="cuda") - - out = OPS[op_name](x) - expected = _reference(op_name, x) - atol, rtol = _tolerances(dtype) - torch.testing.assert_close(out, expected, atol=atol, rtol=rtol) - - -@pytest.mark.parametrize("op_name", OPS) -@pytest.mark.parametrize("dtype", DTYPES) -def test_activation_out_param(op_name: str, dtype: torch.dtype) -> None: - torch.manual_seed(0) - x = torch.randn((4, 7, 128), dtype=dtype, device="cuda") - out = torch.empty((4, 7, 64), dtype=dtype, device="cuda") - - result = OPS[op_name](x, out) - assert result is out - - expected = _reference(op_name, x) - atol, rtol = _tolerances(dtype) - torch.testing.assert_close(out, expected, atol=atol, rtol=rtol) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/python/sglang/multimodal_gen/runtime/layers/activation.py b/python/sglang/multimodal_gen/runtime/layers/activation.py index 1e9e9fec3..b4d457272 100644 --- a/python/sglang/multimodal_gen/runtime/layers/activation.py +++ b/python/sglang/multimodal_gen/runtime/layers/activation.py @@ -16,14 +16,8 @@ from sglang.multimodal_gen.runtime.platforms import current_platform _is_cuda = current_platform.is_cuda() _is_hip = current_platform.is_hip() _is_npu = current_platform.is_npu() -if _is_cuda: - from sglang.jit_kernel.activation import ( - gelu_and_mul, - gelu_tanh_and_mul, - silu_and_mul, - ) -elif _is_hip: - from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul +if _is_cuda or _is_hip: + from sgl_kernel import silu_and_mul if _is_npu: import torch_npu @@ -82,16 +76,8 @@ class GeluAndMul(CustomOp): if approximate not in ("none", "tanh"): raise ValueError(f"Unknown approximate mode: {approximate}") - def forward_cuda(self, x: torch.Tensor) -> Any: - d = x.shape[-1] // 2 - output_shape = x.shape[:-1] + (d,) - out = torch.empty(output_shape, dtype=x.dtype, device=x.device) - if self.approximate == "tanh": - gelu_and_mul_fn = gelu_tanh_and_mul - else: - gelu_and_mul_fn = gelu_and_mul - gelu_and_mul_fn(x, out) - return out + def forward_cuda(self, *args, **kwargs) -> Any: + return self.forward_native(*args, **kwargs) def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 8017ac668..a642fea95 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -49,13 +49,7 @@ _is_cpu = is_cpu() _is_hip = is_hip() _is_xpu = is_xpu() -if _is_cuda: - from sglang.jit_kernel.activation import ( - gelu_and_mul, - gelu_tanh_and_mul, - silu_and_mul, - ) -elif _is_xpu: +if _is_cuda or _is_xpu: from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul elif _is_hip: from sgl_kernel import gelu_and_mul, gelu_quick, gelu_tanh_and_mul, silu_and_mul diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index 3716fc6a6..198599e43 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -17,9 +17,9 @@ if _is_cuda: fp8_blockwise_scaled_grouped_mm, prepare_moe_input, shuffle_rows, + silu_and_mul, ) - from sglang.jit_kernel.activation import silu_and_mul from sglang.jit_kernel.nvfp4 import ( cutlass_fp4_group_mm, scaled_fp4_experts_quant, diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index f06878001..6ce5a3ddc 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -5,9 +5,8 @@ from typing import Optional import torch -from sglang.srt.utils import is_cuda, is_cuda_alike +from sglang.srt.utils import is_cuda_alike -_is_cuda = is_cuda() _is_cuda_alike = is_cuda_alike() if _is_cuda_alike: @@ -16,10 +15,7 @@ if _is_cuda_alike: get_cutlass_w4a8_moe_mm_data, ) -if _is_cuda: - from sglang.jit_kernel.activation import silu_and_mul -else: - from sgl_kernel import silu_and_mul +from sgl_kernel import silu_and_mul from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 from sglang.srt.distributed import get_moe_expert_parallel_world_size diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index 23184e94b..4410f07f3 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py @@ -8,9 +8,8 @@ from sglang.srt.utils.custom_op import register_custom_op _is_cuda = is_cuda() if _is_cuda: - from sgl_kernel import moe_sum_reduce + from sgl_kernel import moe_sum_reduce, silu_and_mul - from sglang.jit_kernel.activation import silu_and_mul from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index 059bc3892..f2621473c 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -48,9 +48,7 @@ _use_sgl_xpu = use_intel_xpu_backend() from sglang.srt.server_args import get_global_server_args if _is_cuda: - from sgl_kernel import moe_sum_reduce - - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul + from sgl_kernel import gelu_and_mul, moe_sum_reduce, silu_and_mul elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py index 15ce91716..b7e6cd7b5 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py @@ -5,6 +5,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Optional import torch +from sgl_kernel import gelu_and_mul, silu_and_mul from triton_kernels.matmul_ogs import ( FlexCtx, FnSpecs, @@ -16,13 +17,6 @@ from triton_kernels.numerics import InFlexData from triton_kernels.routing import GatherIndx, RoutingData, ScatterIndx from triton_kernels.swiglu import swiglu_fn -from sglang.srt.utils import is_cuda - -if is_cuda(): - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul -else: - from sgl_kernel import gelu_and_mul, silu_and_mul - if TYPE_CHECKING: from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.topk import TopKOutput diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index f42b6de6d..93bb9cbcb 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -44,7 +44,7 @@ _is_cuda = is_cuda() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if not (_is_npu or _is_hip) and _is_cuda: - from sglang.jit_kernel.activation import silu_and_mul + from sgl_kernel import silu_and_mul _MASKED_GEMM_FAST_ACT = get_bool_env_var("SGLANG_MASKED_GEMM_FAST_ACT") diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index 5496f7620..be40253b3 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -37,25 +37,26 @@ _is_xpu = is_xpu() _MOE_PADDING_SIZE = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 -if _is_cuda: - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul -elif _is_hip: +if _is_cuda or _is_hip: from sgl_kernel import gelu_and_mul, silu_and_mul - _has_vllm = False - if _use_aiter: - try: - from aiter import moe_sum - except ImportError: - raise ImportError("aiter is required when SGLANG_USE_AITER is set to True") - else: - try: - from vllm import _custom_ops as vllm_ops # moe_sum + if _is_hip: + _has_vllm = False + if _use_aiter: + try: + from aiter import moe_sum + except ImportError: + raise ImportError( + "aiter is required when SGLANG_USE_AITER is set to True" + ) + else: + try: + from vllm import _custom_ops as vllm_ops # moe_sum - _has_vllm = True - except ImportError: - # Fallback: vllm not available, will use triton moe_sum - _has_vllm = False + _has_vllm = True + except ImportError: + # Fallback: vllm not available, will use triton moe_sum + _has_vllm = False elif _is_cpu and _is_cpu_amx_available: pass elif _is_xpu: diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index 3bc7b18a7..748b3c294 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -33,19 +33,7 @@ _is_hip = is_hip() _is_xpu = is_xpu() _is_musa = is_musa() -if _is_cuda: - from sgl_kernel import moe_align_block_size, moe_sum - from sgl_kernel.quantization import ( - ggml_dequantize, - ggml_moe_a8, - ggml_moe_a8_vec, - ggml_moe_get_block_size, - ggml_mul_mat_a8, - ggml_mul_mat_vec_a8, - ) - - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul -elif _is_musa: +if _is_cuda or _is_musa: from sgl_kernel import gelu_and_mul, moe_align_block_size, moe_sum, silu_and_mul from sgl_kernel.quantization import ( ggml_dequantize, @@ -200,11 +188,16 @@ def fused_moe_gguf( activation: str, ) -> torch.Tensor: def act(x: torch.Tensor): + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) if activation == "silu": - return silu_and_mul(x) + silu_and_mul(out, x) elif activation == "gelu": - return gelu_and_mul(x) - raise ValueError(f"Unsupported activation: {activation}") + gelu_and_mul(out, x) + else: + raise ValueError(f"Unsupported activation: {activation}") + return out out_hidden_states = torch.empty_like(x) # unless we decent expert reuse we are better off running moe_vec kernel diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index ed835bcb4..3060e2fad 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -47,11 +47,12 @@ _use_aiter = bool(int(os.getenv("SGLANG_USE_AITER", "0"))) _is_xpu = is_xpu() _MOE_PADDING_SIZE = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 -if _is_cuda: - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul -elif _is_hip: + +if _is_cuda or _is_hip: from sgl_kernel import gelu_and_mul, silu_and_mul - from vllm import _custom_ops as vllm_ops # moe_sum + + if _is_hip: + from vllm import _custom_ops as vllm_ops # moe_sum elif _is_cpu and _is_cpu_amx_available: pass elif _is_xpu: