diff --git a/.claude/skills/generate-profile/SKILL.md b/.claude/skills/generate-profile/SKILL.md new file mode 100644 index 000000000..2b6201f6b --- /dev/null +++ b/.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= sglang serve --model-path --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:/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//` +4. The trace directory contains: + - `-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/ +``` + +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 ` — 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 ` and use `--host 127.0.0.1 --port ` for test commands +- **Multi-GPU**: Use `--tp ` 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. diff --git a/python/sglang/jit_kernel/benchmark/bench_clamp_position.py b/python/sglang/jit_kernel/benchmark/bench_clamp_position.py new file mode 100644 index 000000000..08fa92660 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/clamp_position.py b/python/sglang/jit_kernel/clamp_position.py new file mode 100644 index 000000000..ed57da776 --- /dev/null +++ b/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 diff --git a/python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh b/python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh new file mode 100644 index 000000000..963dcba25 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh @@ -0,0 +1,54 @@ +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For div_ceil + +#include // For LaunchKernel + +#include +#include + +#include +#include + +namespace { + +template +__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 +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(); + + TensorMatcher({N}) // + .with_dtype() + .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, + static_cast(dst.data_ptr()), + static_cast(seq_lens.data_ptr()), + num_elements); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/tests/test_clamp_position.py b/python/sglang/jit_kernel/tests/test_clamp_position.py new file mode 100644 index 000000000..4cabc7a4b --- /dev/null +++ b/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"]) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 66b144f61..6c58d0a80 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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 diff --git a/python/sglang/test/send_one.py b/python/sglang/test/send_one.py index 46838a44a..260019abc 100644 --- a/python/sglang/test/send_one.py +++ b/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