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

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

.claude/skills/generate-profile/SKILL.md

@@ -0,0 +1,141 @@
+---
+name: generate-profile
+description: Generate an e2e profiling trace of an SGLang server run. Launches a server, validates accuracy, captures a Chrome-compatible trace, and returns the profile path.
+---
+
+# Generate an E2E Profile of an SGLang Server Run
+
+This skill launches an SGLang server, validates it with a quick accuracy test, generates a profiling trace, and returns the profile file path.
+
+## Prerequisites
+
+- A working SGLang installation (`pip install -e .` or equivalent)
+- At least one available CUDA GPU
+
+## Step-by-step Workflow
+
+### Step 1: Launch the server
+
+```bash
+CUDA_VISIBLE_DEVICES=<gpu_id> sglang serve --model-path <model> --port <port> &
+```
+
+- Default model: `Qwen/Qwen3-8B` (good balance of speed and quality)
+- Default port: `30000`
+- The server runs in the background. Save the PID for cleanup.
+- Use the GPU specified by the user's preferences (check memory files for GPU preferences).
+
+### Step 2: Wait for server readiness
+
+Poll the health endpoint until the server is ready:
+
+```bash
+for i in $(seq 1 120); do
+  if curl -s http://127.0.0.1:<port>/health 2>/dev/null | grep -q "ok\|healthy"; then
+    echo "Server ready"
+    break
+  fi
+  sleep 5
+done
+```
+
+The server prints **"The server is fired up and ready to roll!"** to stdout when ready. The health endpoint returns 200 once the server can accept requests.
+
+Typical startup time: 30-90 seconds depending on model size and whether CUDA graphs are being compiled.
+
+### Step 3: Validate accuracy (sanity check)
+
+```bash
+python3 -m sglang.test.few_shot_gsm8k --num-q 20
+```
+
+- Expected accuracy: **> 0.8** for capable models (Qwen3-8B, Llama-3.1-8B-Instruct, etc.)
+- This is a quick sanity check, not a rigorous benchmark.
+- If accuracy is unexpectedly low, something is wrong — do not proceed to profiling.
+
+### Step 4: Generate the profile
+
+```bash
+python3 -m sglang.test.send_one --profile
+```
+
+This command:
+1. Sends a request to the server
+2. Triggers the profiler for 5 steps (default)
+3. Generates a trace file under `/tmp/<timestamp>/`
+4. The trace directory contains:
+   - `<timestamp>-TP-0.trace.json.gz` — Chrome trace format (open in `chrome://tracing` or Perfetto)
+   - `server_args.json` — the server configuration used
+
+**Output format:**
+```
+Dump profiling traces to /tmp/<timestamp>
+```
+
+The profile path is printed to stdout. Parse it from the output.
+
+**Optional flags:**
+- `--profile-steps N` — number of profiling steps (default: 5)
+- `--profile-by-stage` — profile by stage (prefill/decode separately)
+- `--profile-prefix <path>` — custom output prefix
+
+### Step 5: Kill the server
+
+```bash
+pkill -9 -f "sglang.launch_server\|sglang serve\|sglang.srt"
+```
+
+Wait a moment and verify no sglang processes remain:
+```bash
+sleep 2 && pgrep -af "sglang serve" || echo "Server killed"
+```
+
+### Step 6: Report the profile path
+
+Return the profile directory path (e.g., `/tmp/1773999986.4769795`) and list its contents so the user knows what files were generated.
+
+## Example Full Run
+
+```bash
+# 1. Launch server
+source cleanup/bin/activate
+CUDA_VISIBLE_DEVICES=1 sglang serve --model-path Qwen/Qwen3-8B --port 30000 &
+
+# 2. Wait for ready
+for i in $(seq 1 120); do
+  curl -s http://127.0.0.1:30000/health | grep -q "ok" && break
+  sleep 5
+done
+
+# 3. Accuracy check
+python3 -m sglang.test.few_shot_gsm8k --num-q 20
+# Expected: Accuracy > 0.8
+
+# 4. Profile
+python3 -m sglang.test.send_one --profile
+# Output: "Dump profiling traces to /tmp/1773999986.4769795"
+
+# 5. Cleanup
+pkill -9 -f "sglang.launch_server\|sglang serve\|sglang.srt"
+sleep 2
+
+# 6. Check output
+ls -la /tmp/1773999986.4769795/
+# 1773999986.4851577-TP-0.trace.json.gz  (Chrome trace)
+# server_args.json                        (server config)
+```
+
+## Customization
+
+- **Different port**: Pass `--port <port>` and use `--host 127.0.0.1 --port <port>` for test commands
+- **Multi-GPU**: Use `--tp <N>` for tensor parallelism; trace files will be generated per TP rank
+- **Longer profile**: Use `--profile-steps 10` for more steps in the trace
+- **Stage profiling**: Use `--profile-by-stage` to separate prefill and decode phases
+
+## Viewing the Profile
+
+Open the `.trace.json.gz` file in:
+- **Perfetto UI**: https://ui.perfetto.dev/ (drag and drop the file)
+- **Chrome tracing**: `chrome://tracing` (load the file)
+
+Both support the gzipped Chrome trace format natively.

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