Replace clamp_position with JIT kernel + platform dispatch (#20999)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

python/sglang/jit_kernel/benchmark/bench_clamp_position.py

@@ -0,0 +1,61 @@
+import itertools
+
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.benchmark.utils import (
+    DEFAULT_DEVICE,
+    get_benchmark_range,
+    run_benchmark,
+)
+from sglang.jit_kernel.clamp_position import clamp_position_cuda
+from sglang.srt.utils import get_compiler_backend
+
+SIZE_LIST = get_benchmark_range(
+    full_range=[2**n for n in range(4, 16)],
+    ci_range=[256, 4096],
+)
+
+configs = list(itertools.product(SIZE_LIST))
+
+
+def _torch_clamp_position(seq_lens):
+    return torch.clamp(seq_lens - 1, min=0).to(torch.int64)
+
+
+_compiled_clamp_position = torch.compile(
+    _torch_clamp_position, dynamic=True, backend=get_compiler_backend()
+)
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["size"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["jit", "torch_compile", "torch"],
+        line_names=["SGL JIT Kernel", "torch.compile", "PyTorch"],
+        styles=[("blue", "-"), ("green", "-."), ("red", "--")],
+        ylabel="us",
+        plot_name="clamp-position-performance",
+        args={},
+    )
+)
+def benchmark(size: int, provider: str):
+    seq_lens = torch.randint(
+        0, 10000, (size,), dtype=torch.int64, device=DEFAULT_DEVICE
+    )
+
+    if provider == "jit":
+        fn = lambda: clamp_position_cuda(seq_lens)
+    elif provider == "torch_compile":
+        fn = lambda: _compiled_clamp_position(seq_lens)
+    else:
+        fn = lambda: _torch_clamp_position(seq_lens)
+
+    return run_benchmark(fn)
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)

python/sglang/jit_kernel/clamp_position.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_clamp_position_module(dtype: torch.dtype) -> Module:
+    """Compile and cache the JIT clamp_position module for a given dtype."""
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "clamp_position",
+        *args,
+        cuda_files=["elementwise/clamp_position.cuh"],
+        cuda_wrappers=[
+            ("clamp_position", f"ClampPosition<{args}>::run"),
+        ],
+    )
+
+
+def clamp_position_cuda(seq_lens: torch.Tensor) -> torch.Tensor:
+    """Compute positions = clamp(seq_lens - 1, min=0) on CUDA.
+
+    Supported dtypes: torch.int32, torch.int64.
+    """
+    dst = torch.empty_like(seq_lens)
+    module = _jit_clamp_position_module(seq_lens.dtype)
+    module.clamp_position(dst, seq_lens)
+    return dst

python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh

@@ -0,0 +1,54 @@
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+template <typename T>
+__global__ void clamp_position_kernel(T* __restrict__ dst, const T* __restrict__ seq_lens, size_t n) {
+  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    T val = seq_lens[idx] - 1;
+    dst[idx] = val < 0 ? 0 : val;
+  }
+}
+
+constexpr size_t kBlockSize = 256;
+
+template <typename T>
+struct ClampPosition {
+  static void run(tvm::ffi::TensorView dst, tvm::ffi::TensorView seq_lens) {
+    using namespace host;
+
+    SymbolicSize N = {"num_elements"};
+    SymbolicDevice device_;
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .verify(dst)
+        .verify(seq_lens);
+
+    const size_t num_elements = N.unwrap();
+    if (num_elements == 0) return;
+
+    const size_t grid_size = div_ceil(num_elements, kBlockSize);
+    const DLDevice device = device_.unwrap();
+
+    LaunchKernel(grid_size, kBlockSize, device)(
+        clamp_position_kernel<T>,
+        static_cast<T*>(dst.data_ptr()),
+        static_cast<const T*>(seq_lens.data_ptr()),
+        num_elements);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/tests/test_clamp_position.py

@@ -0,0 +1,40 @@
+import pytest
+import torch
+
+from sglang.jit_kernel.clamp_position import clamp_position_cuda
+
+
+def _reference_clamp_position(seq_lens):
+    return torch.clamp(seq_lens - 1, min=0).to(seq_lens.dtype)
+
+
+@pytest.mark.parametrize("size", [1, 2, 127, 128, 255, 256, 1024, 4097])
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+class TestClampPosition:
+    def test_normal(self, size: int, dtype: torch.dtype) -> None:
+        seq_lens = torch.randint(1, 10000, (size,), dtype=dtype, device="cuda")
+        expected = _reference_clamp_position(seq_lens)
+        result = clamp_position_cuda(seq_lens)
+        assert torch.equal(result, expected)
+
+    def test_zeros(self, size: int, dtype: torch.dtype) -> None:
+        seq_lens = torch.zeros(size, dtype=dtype, device="cuda")
+        expected = _reference_clamp_position(seq_lens)
+        result = clamp_position_cuda(seq_lens)
+        assert torch.equal(result, expected)
+
+    def test_ones(self, size: int, dtype: torch.dtype) -> None:
+        seq_lens = torch.ones(size, dtype=dtype, device="cuda")
+        expected = _reference_clamp_position(seq_lens)
+        result = clamp_position_cuda(seq_lens)
+        assert torch.equal(result, expected)
+
+    def test_mixed(self, size: int, dtype: torch.dtype) -> None:
+        seq_lens = torch.randint(0, 10000, (size,), dtype=dtype, device="cuda")
+        expected = _reference_clamp_position(seq_lens)
+        result = clamp_position_cuda(seq_lens)
+        assert torch.equal(result, expected)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

python/sglang/srt/model_executor/forward_batch_info.py

@@ -56,7 +56,13 @@ from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
     ForwardBatchDeepSeekMHAMixin,
 )
 from sglang.srt.server_args import get_global_server_args
-from sglang.srt.utils import get_compiler_backend, is_hip, is_npu, support_triton
+from sglang.srt.utils import (
+    get_compiler_backend,
+    is_cuda,
+    is_hip,
+    is_npu,
+    support_triton,
+)
 from sglang.srt.utils.common import ceil_align
 
 if TYPE_CHECKING:
@@ -1175,6 +1181,17 @@ def compute_position_torch(
     return positions.to(torch.int64), extend_start_loc
 
 
-@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
-def clamp_position(seq_lens):
+def _clamp_position_native(seq_lens):
     return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
+
+
+if is_cuda():
+    from sglang.jit_kernel.clamp_position import clamp_position_cuda
+
+    clamp_position = clamp_position_cuda
+elif is_hip():
+    clamp_position = torch.compile(
+        _clamp_position_native, dynamic=True, backend=get_compiler_backend()
+    )
+else:
+    clamp_position = _clamp_position_native

python/sglang/test/send_one.py

@@ -40,7 +40,7 @@ class BenchArgs:
     stop: Optional[list] = None
     stream: bool = False
     profile: bool = False
-    profile_steps: int = 3
+    profile_steps: int = 5
     profile_by_stage: bool = False
     profile_prefix: Optional[str] = None