[Docs] Sync docs_new with legacy docs and update migration redirects (#23337)
Co-authored-by: Mingyi <wisclmy0611@gmail.com>
This commit is contained in:
zijiexia
@@ -142,7 +142,7 @@ python3 -m sglang.bench_serving \
|
||||
|
||||
- `--output-file FILE.jsonl`: append JSONL results to file; auto-named if unspecified
|
||||
- `--output-details`: include per-request arrays (generated texts, errors, ttfts, itls, input/output lens)
|
||||
- `--extra-request-body '{"top_p":0.9,"temperature":0.6}'`: merged into payload (sampling params, etc.)
|
||||
- `--extra-request-body '{"top_p":0.9,"temperature":0.6}'`: merged into payload (sampling params, etc.)
|
||||
- `--disable-ignore-eos`: pass through EOS behavior (varies by backend)
|
||||
- `--warmup-requests N`: run warmup requests with short output first (default 1)
|
||||
- `--flush-cache`: call `/flush_cache` (sglang) before main run
|
||||
@@ -335,7 +335,7 @@ python3 -m sglang.bench_serving \
|
||||
python3 -m sglang.bench_serving \
|
||||
--backend sglang \
|
||||
--host 127.0.0.1 --port 30000 \
|
||||
--model mode-name \
|
||||
--model model-name \
|
||||
--dataset-name mooncake \
|
||||
--mooncake-slowdown-factor 1.0 \
|
||||
--mooncake-num-rounds 1000 \
|
||||
@@ -344,6 +344,41 @@ python3 -m sglang.bench_serving \
|
||||
--random-output-len 256
|
||||
```
|
||||
|
||||
10) Fake decode stress testing (PD disaggregation, decode-only):
|
||||
|
||||
When benchmarking pure decode performance in a PD disaggregation setup, you can bypass the prefill node entirely by using `--fake-prefill`. This requires the decode server to be started with `--disaggregation-transfer-backend fake`:
|
||||
|
||||
```bash Command
|
||||
# Step 1: Start a decode-only server with fake transfer backend
|
||||
python -m sglang.launch_server \
|
||||
--model-path meta-llama/Llama-3.1-8B-Instruct \
|
||||
--disaggregation-mode decode \
|
||||
--disaggregation-transfer-backend fake \
|
||||
--port 30001
|
||||
|
||||
# Step 2: Run bench_serving with --fake-prefill
|
||||
python3 -m sglang.bench_serving \
|
||||
--backend sglang \
|
||||
--host 127.0.0.1 --port 30001 \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--dataset-name random \
|
||||
--num-prompts 500 \
|
||||
--random-input-len 1024 --random-output-len 256 \
|
||||
--fake-prefill
|
||||
```
|
||||
|
||||
Similarly, `bench_one_batch_server` also supports `--fake-prefill`:
|
||||
|
||||
```bash Command
|
||||
python3 -m sglang.bench_one_batch_server \
|
||||
--base-url http://127.0.0.1:30001 \
|
||||
--model-path meta-llama/Llama-3.1-8B-Instruct \
|
||||
--batch-size 32 --input-len 1024 --output-len 256 \
|
||||
--fake-prefill
|
||||
```
|
||||
|
||||
The `--fake-prefill` flag automatically injects special sentinel values into each request, telling the decode server to skip real KV transfer and generate fake KV data locally.
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
- All requests failed: verify `--backend`, server URL/port, `--model`, and authentication. Check warmup errors printed by the script.
|
||||
@@ -355,4 +390,4 @@ python3 -m sglang.bench_serving \
|
||||
### Notes
|
||||
|
||||
- The script raises the file descriptor soft limit (`RLIMIT_NOFILE`) to help with many concurrent connections.
|
||||
- For sglang, `/get_server_info` is queried post-run to report speculative decoding accept length when available.
|
||||
- For sglang, `/server_info` is queried post-run to report speculative decoding accept length when available.
|
||||
|
||||
@@ -5,28 +5,69 @@ metatags:
|
||||
---
|
||||
## Benchmark
|
||||
|
||||
- Benchmark the latency of running a single static batch without a server. The arguments are the same as for `launch_server.py`.
|
||||
Note that this is a simplified test script without a dynamic batching server, so it may run out of memory for a batch size that a real server can handle. A real server truncates the prefill into several batches, while this simplified script does not.
|
||||
- Without a server (do not need to launch a server)
|
||||
```bash Command
|
||||
python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch 32 --input-len 256 --output-len 32
|
||||
```
|
||||
- With a server (please use `sglang.launch_server` to launch a server first and run the following command.)
|
||||
```bash Command
|
||||
python -m sglang.bench_one_batch_server --base-url http://127.0.0.1:30000 --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch-size 32 --input-len 256 --output-len 32
|
||||
```
|
||||
SGLang provides four benchmark tools that operate at different levels of the stack. The table below summarizes their key differences:
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Tool</th>
|
||||
<th>HTTP Server</th>
|
||||
<th>Scheduler</th>
|
||||
<th>Use Case</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td><code>bench_serving</code></td>
|
||||
<td>Yes (async HTTP client to a running server)</td>
|
||||
<td>Yes (indirectly, via server)</td>
|
||||
<td>Realistic online serving benchmarks with latency metrics (TTFT, TPOT, ITL)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>bench_one_batch_server</code></td>
|
||||
<td>Yes (sends HTTP requests to a running server)</td>
|
||||
<td>Yes (indirectly, via server)</td>
|
||||
<td>End-to-end single-batch latency including HTTP and scheduler overhead</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>bench_offline_throughput</code></td>
|
||||
<td>No</td>
|
||||
<td>Yes (directly uses <code>Engine</code> in-process)</td>
|
||||
<td>Maximum throughput measurement without HTTP overhead</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>bench_one_batch</code></td>
|
||||
<td>No</td>
|
||||
<td>No (directly calls <code>ModelRunner</code>)</td>
|
||||
<td>Kernel-level latency profiling of a single static batch</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
- Benchmark offline processing. This script will start an offline engine and run the benchmark.
|
||||
Use `bench_serving` by default unless there are specific needs.
|
||||
|
||||
**`bench_serving`** is an async HTTP load-testing client that sends requests at controlled rates with configurable concurrency to a running server. It measures realistic online serving metrics including time-to-first-token (TTFT), time-per-output-token (TPOT), inter-token latency (ITL), and throughput. Use `num-prompts >= 5 * max-concurrency` to measure steady-state performance. Launch a server with `sglang.launch_server` first.
|
||||
|
||||
```bash Command
|
||||
python3 -m sglang.bench_serving --backend sglang --max-concurrency 16 --num-prompts 80 --random-input-len 256 --random-output-len 32 --dataset-name random
|
||||
```
|
||||
|
||||
**`bench_one_batch_server`** sends a single batch as one HTTP request to a running server. Due to only having a single batch, the server is never in a steady-state and metrics will be biased. Launch a server with `sglang.launch_server` first.
|
||||
|
||||
```bash Command
|
||||
python3 -m sglang.bench_one_batch_server --base-url http://127.0.0.1:30000 --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch-size 32 --input-len 256 --output-len 32
|
||||
```
|
||||
|
||||
**`bench_offline_throughput`** directly instantiates the `Engine` object in-process (no HTTP server) and submits all requests at once via `engine.generate()`. The engine's scheduler handles batching and execution. This measures maximum achievable throughput without any network overhead.
|
||||
|
||||
```bash Command
|
||||
python3 -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10
|
||||
```
|
||||
|
||||
- Benchmark online serving. Please use `sglang.launch_server` to launch a server first and run the following command.
|
||||
**`bench_one_batch`** is the lowest-level tool. It directly instantiates a `ModelRunner` and calls `extend()` / `decode()` on a fixed static batch, bypassing the scheduler entirely. The prefill and decode phases are run separately, making profiling easier but rendering the metrics unrealistic. Because there is no dynamic batching, it may run out of memory for batch sizes that a real server can handle (a real server chunks prefill into smaller batches). This is best suited for profiling individual kernel performance.
|
||||
|
||||
```bash Command
|
||||
python3 -m sglang.bench_serving --backend sglang --num-prompt 10
|
||||
python3 -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch-size 32 --input-len 256 --output-len 32
|
||||
```
|
||||
|
||||
## Profile with PyTorch Profiler
|
||||
@@ -46,7 +87,10 @@ python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct
|
||||
python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --sharegpt-output-len 100 --profile
|
||||
```
|
||||
|
||||
Please make sure that the `SGLANG_TORCH_PROFILER_DIR` should be set at both server and client side, otherwise the trace file cannot be generated correctly . A secure way will be setting `SGLANG_TORCH_PROFILER_DIR` in the `.*rc` file of shell (e.g. `~/.bashrc` for bash shells).
|
||||
For `bench_serving --profile`, the output directory is selected on the client side from `--profile-output-dir` or `SGLANG_TORCH_PROFILER_DIR` (fallback: `/tmp`), then sent in the `/start_profile` request.
|
||||
If you call `/start_profile` directly and do not provide `output_dir`, the server uses its own `SGLANG_TORCH_PROFILER_DIR` (fallback: `/tmp`).
|
||||
|
||||
Setting `SGLANG_TORCH_PROFILER_DIR` on both server and client is still recommended to avoid confusion about where traces are written.
|
||||
|
||||
For more details, please refer to [Bench Serving Guide](./bench_serving).
|
||||
|
||||
@@ -147,7 +191,7 @@ curl -X POST http://127.0.0.1:30000/start_profile \
|
||||
**Parameters:**
|
||||
|
||||
- `output_dir` (optional): Directory where profile traces will be saved. If not specified, uses `SGLANG_TORCH_PROFILER_DIR` environment variable, or `/tmp` as the default
|
||||
- `num_steps` (optional): Number of steps to profile. If not specified, profiling continues until manually stopped with `/end_profile`
|
||||
- `num_steps` (optional): Number of steps to profile. If not specified, profiling continues until manually stopped with `/stop_profile`
|
||||
- `start_step` (optional): Step number at which to start profiling (inclusive). Useful for skipping warmup iterations
|
||||
- `activities` (optional): List of activities to profile, e.g., `["CPU", "GPU"]`. Default is `["CPU", "GPU"]`
|
||||
- `merge_profiles` (optional): Whether to merge distributed traces. Default is `false`
|
||||
@@ -171,17 +215,17 @@ curl -X POST http://127.0.0.1:30000/start_profile \
|
||||
**Continuous profiling (manual stop):**
|
||||
|
||||
```bash Command
|
||||
# Start profiling without num_steps - must manually stop with /end_profile
|
||||
# Start profiling without num_steps - must manually stop with /stop_profile
|
||||
curl -X POST http://127.0.0.1:30000/start_profile
|
||||
```
|
||||
|
||||
#### Using `/end_profile` endpoint
|
||||
#### Using `/stop_profile` endpoint
|
||||
|
||||
The `/end_profile` endpoint stops an ongoing profiling session and saves the trace file.
|
||||
The `/stop_profile` endpoint stops an ongoing profiling session and saves the trace file.
|
||||
|
||||
```bash Command
|
||||
# Stop profiling and save traces
|
||||
curl -X POST http://127.0.0.1:30000/end_profile
|
||||
curl -X POST http://127.0.0.1:30000/stop_profile
|
||||
```
|
||||
|
||||
This is only needed when you start profiling without specifying `num_steps`. If `num_steps` is specified, profiling will automatically stop after that many steps.
|
||||
@@ -204,7 +248,7 @@ curl -X POST http://127.0.0.1:30000/start_profile \
|
||||
python -m sglang.bench_serving --backend sglang --num-prompts 100
|
||||
|
||||
# Terminal 2: Stop profiling when done
|
||||
curl -X POST http://127.0.0.1:30000/end_profile
|
||||
curl -X POST http://127.0.0.1:30000/stop_profile
|
||||
```
|
||||
|
||||
### Profiler Trace Merger for Distributed Traces
|
||||
@@ -246,8 +290,8 @@ python -m sglang.profiler \
|
||||
#### Output Files
|
||||
|
||||
The profile merger generates:
|
||||
- Individual rank trace files: `{profile_id}-TP-{tp}-DP-{dp}-PP-{pp}-EP-{ep}.trace.json.gz`
|
||||
- Merged trace file: `merged-{profile_id}.trace.json.gz`
|
||||
- Individual rank trace files: `{profile_id}-TP-{tp}-DP-{dp}-PP-{pp}-EP-{ep}.trace.json.gz`
|
||||
- Merged trace file: `merged-{profile_id}.trace.json.gz`
|
||||
|
||||
### Possible PyTorch bugs
|
||||
If in any cases you encounter the following error (for example, using qwen 2.5 VL):
|
||||
@@ -398,10 +442,10 @@ This method allows you to control exactly when profiling starts/stops via HTTP A
|
||||
|
||||
```bash Command
|
||||
# Terminal 2: Only needed if num_steps was not specified
|
||||
curl -X POST http://127.0.0.1:30000/end_profile
|
||||
curl -X POST http://127.0.0.1:30000/stop_profile
|
||||
```
|
||||
|
||||
The `--capture-range=cudaProfilerApi` option tells Nsight Systems to only capture data between `cudaProfilerStart()` and `cudaProfilerStop()` calls (triggered by `/start_profile` and `/end_profile`), reducing overhead and file size. The `start_step` parameter skips the first 3 steps to avoid capturing warmup overhead.
|
||||
The `--capture-range=cudaProfilerApi` option tells Nsight Systems to only capture data between `cudaProfilerStart()` and `cudaProfilerStop()` calls (triggered by `/start_profile` and `/stop_profile`), reducing overhead and file size. The `start_step` parameter skips the first 3 steps to avoid capturing warmup overhead.
|
||||
|
||||
**Method 2: Simpler approach without `/start_profile` API**
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ mode: wide
|
||||
metatags:
|
||||
description: "SGLang contribution guide: source install, pre-commit, unit tests, CI triggers, code style, sgl-kernel updates."
|
||||
---
|
||||
Welcome to **SGLang**! We appreciate your interest in contributing. This guide provides a concise overview of how to set up your environment, run tests, build documentation, and open a Pull Request (PR). Whether you're fixing a small bug or developing a major feature, we encourage following these steps for a smooth contribution process.
|
||||
Welcome to **SGLang**! We appreciate your interest in contributing. This guide provides a concise overview of how to set up your environment, run tests, build documentation, and open a Pull Request (PR). Whether you’re fixing a small bug or developing a major feature, we encourage following these steps for a smooth contribution process.
|
||||
|
||||
## Install SGLang from Source
|
||||
|
||||
@@ -18,7 +18,7 @@ git clone https://github.com/<your_user_name>/sglang.git
|
||||
|
||||
### Build from source
|
||||
|
||||
Refer to [Install SGLang from Source](../get-started/installation).
|
||||
Refer to [Install SGLang from Source](../get-started/install#method-2-from-source).
|
||||
|
||||
## Format code with pre-commit
|
||||
|
||||
@@ -32,17 +32,50 @@ pre-commit run --all-files
|
||||
|
||||
- **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request.
|
||||
- **Do not commit** directly to the `main` branch. Always create a new branch (e.g., `feature/my-new-feature`), push your changes, and open a PR from that branch.
|
||||
- Link checking with lychee is **enforced in CI**. By default, it is not blocking local commits.
|
||||
- To run local link checks manually, use: `pre-commit run --hook-stage manual lychee --all-files`.
|
||||
|
||||
## Run and add unit tests
|
||||
|
||||
If you add a new feature or fix a bug, please add corresponding unit tests to ensure coverage and prevent regression.
|
||||
SGLang uses Python's built-in [unittest](https://docs.python.org/3/library/unittest.html) framework.
|
||||
For detailed instructions on running tests and integrating them into CI, refer to [test/README](https://github.com/sgl-project/sglang/tree/main/test/README).
|
||||
|
||||
### Unit tests (no server required)
|
||||
|
||||
Unit tests live under [`test/registered/unit/`](https://github.com/sgl-project/sglang/tree/main/test/registered/unit), organized to mirror the `python/sglang/srt/` source tree. These tests validate component logic **without** launching a server or loading real model weights.
|
||||
SGLang uses Python's built-in [unittest](https://docs.python.org/3/library/unittest.html) framework with [pytest](https://docs.pytest.org/) as the test runner.
|
||||
|
||||
**When to add a unit test:** If you modify a file under `python/sglang/srt/`, check whether a corresponding test exists in `test/registered/unit/` and add coverage for your changes. For example:
|
||||
|
||||
```
|
||||
srt/mem_cache/radix_cache.py → unit/mem_cache/test_radix_cache.py
|
||||
srt/sampling/sampling_params.py → unit/sampling/test_sampling_params.py
|
||||
```
|
||||
|
||||
**Run unit tests locally:**
|
||||
|
||||
```bash Command
|
||||
pytest test/registered/unit/ -v # all unit tests
|
||||
pytest test/registered/unit/mem_cache/ -v # one module
|
||||
```
|
||||
|
||||
**Run with coverage:**
|
||||
|
||||
```bash Command
|
||||
pytest test/registered/unit/ --cov --cov-config=.coveragerc -v
|
||||
```
|
||||
|
||||
For conventions on CI registration, test structure, and examples, see [`test/registered/unit/README.md`](https://github.com/sgl-project/sglang/tree/main/test/registered/unit/README.md).
|
||||
|
||||
### E2E tests (server required)
|
||||
|
||||
For tests that require launching a server, refer to [`test/registered/README.md`](https://github.com/sgl-project/sglang/tree/main/test/registered/README.md) for guidance on where to place your test.
|
||||
|
||||
For detailed instructions on running tests and integrating them into CI, refer to [test/README.md](https://github.com/sgl-project/sglang/tree/main/test/README.md).
|
||||
|
||||
## Write documentations
|
||||
|
||||
We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase.
|
||||
For more details, please refer to [docs/README](https://github.com/sgl-project/sglang/blob/main/docs/README.md).
|
||||
For more details, please refer to [docs/README.md](https://github.com/sgl-project/sglang/tree/main/docs/README.md).
|
||||
|
||||
## Test the accuracy
|
||||
If your code changes the model output, please run the accuracy tests. A quick sanity check is the few-shot GSM8K.
|
||||
@@ -56,19 +89,19 @@ python3 -m sglang.test.few_shot_gsm8k --num-questions 200
|
||||
```
|
||||
|
||||
Please note that the above script is primarily a sanity check, not a rigorous accuracy or speed test.
|
||||
This test can have significant variance (1%-5%) in accuracy due to batching and the non-deterministic nature of the inference engine.
|
||||
This test can have significant variance (1%–5%) in accuracy due to batching and the non-deterministic nature of the inference engine.
|
||||
Also, do not rely on the "Latency/Output throughput" from this script, as it is not a proper speed test.
|
||||
|
||||
GSM8K is too easy for state-of-the-art models nowadays. Please try your own more challenging accuracy tests.
|
||||
You can find additional accuracy eval examples in:
|
||||
- [test_eval_accuracy_large.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_eval_accuracy_large.py)
|
||||
- [test_gpt_oss_1gpu.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_gpt_oss_1gpu.py)
|
||||
- [test_eval_accuracy_large.py](https://github.com/sgl-project/sglang/blob/main/test/registered/eval/test_eval_accuracy_large.py)
|
||||
- [test_gpt_oss_1gpu.py](https://github.com/sgl-project/sglang/blob/main/test/registered/core/test_gpt_oss_1gpu.py)
|
||||
|
||||
## Benchmark the speed
|
||||
Refer to [Benchmark and Profiling](../developer_guide/benchmark_and_profiling).
|
||||
Refer to [Benchmark and Profiling](./benchmark_and_profiling).
|
||||
|
||||
## Requesting a review for merge
|
||||
You can follow the pull request merge process described in [MAINTAINER](https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER).
|
||||
You can follow the pull request merge process described in [MAINTAINER.md](https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER.md).
|
||||
You will need to work with the Merge Oncall, Codeowner, and other reviewers to get their approvals.
|
||||
Then your PR can be merged.
|
||||
|
||||
@@ -77,6 +110,8 @@ Then your PR can be merged.
|
||||
We have a lot of open PRs but limited CI machines, so only top and trusted contributors have permission to trigger CI tests.
|
||||
Users with permission are listed in the [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob/main/.github/CI_PERMISSIONS.json)
|
||||
|
||||
**PR authors** can always use `/rerun-failed-ci` on their own PRs, even if they are not listed in `CI_PERMISSIONS.json`.
|
||||
|
||||
For CI to run on a pull request, it must have the "run-ci" label. Authorized users can add the label or rerun failed tests by commenting on the PR with one of these commands:
|
||||
|
||||
- `/tag-run-ci-label`: Adds the "run-ci" label. Every future commit will trigger CI.
|
||||
@@ -84,18 +119,17 @@ For CI to run on a pull request, it must have the "run-ci" label. Authorized use
|
||||
- `/tag-and-rerun-ci`: A single command that performs both `/tag-run-ci-label` and `/rerun-failed-ci`.
|
||||
- `/rerun-stage <stage-name>`: Reruns a specific test stage without waiting for its dependencies. This is useful when you want to quickly validate a fix for a specific test failure instead of waiting ~30 minutes for preceding stages to complete.
|
||||
|
||||
If you have permission, the [Slash Command Handler](https://github.com/sgl-project/sglang/actions/workflows/slash-command-handler.yml) will run your command and react with a +1 to your comment. It may take up to a few minutes for the reaction to appear. Here's a usage [example](https://github.com/sgl-project/sglang/pull/14253#issuecomment-3599509302).
|
||||
If you have permission, the [Slash Command Handler](https://github.com/sgl-project/sglang/actions/workflows/slash-command-handler.yml) will run your command and react with a 👍 to your comment. It may take up to a few minutes for the reaction to appear. Here’s a usage [example](https://github.com/sgl-project/sglang/pull/14253#issuecomment-3599509302).
|
||||
|
||||
To avoid spamming a PR with too many `/rerun-failed-ci` comments, you can also trigger the command by editing an existing comment and adding any suffix (e.g., `/rerun-failed-ci try again`).
|
||||
|
||||
Example of rerunning a single test stage: `/rerun-stage unit-test-backend-4-gpu`.
|
||||
|
||||
If you don't have permission, please ask maintainers to trigger CI for you.
|
||||
If you don’t have permission and you’re not the PR author, please ask maintainers to trigger CI for you.
|
||||
|
||||
### CI rate limits
|
||||
|
||||
Due to CI scheduling and limited resources, higher-priority PRs may preempt running jobs. In such cases, you may need to rerun the tests.
|
||||
|
||||
We apply CI rate limits to prevent abuse and ensure fair usage of our CI resources.
|
||||
|
||||
Each CI workflow has a default limit defined in its workflow configuration file. For example, in [pr-gate.yml](https://github.com/sgl-project/sglang/blob/main/.github/workflows/pr-gate.yml), the default cooldown period is 120 minutes, and each workflow can override it via the `cool-down-minutes` input parameter:
|
||||
@@ -107,8 +141,7 @@ cool-down-minutes:
|
||||
default: 120
|
||||
```
|
||||
|
||||
Users listed in [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob/main/.github/CI_PERMISSIONS.json) may have a per-user cooldown interval. In practice, we use the minimum of the workflow's default window and the user-specific interval.
|
||||
|
||||
Users listed in [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob/main/.github/CI_PERMISSIONS.json) may have a per-user cooldown interval. In practice, we use the minimum of the workflow’s default window and the user-specific interval.
|
||||
|
||||
## Code style guidance
|
||||
- Avoid code duplication. If the same code snippet (more than five lines) appears multiple times, extract it into a shared function.
|
||||
@@ -121,28 +154,34 @@ Users listed in [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob
|
||||
- If a single test file run longer than 500 seconds, split it into multiple smaller files (e.g., `test_eagle_infer_a.py`, `test_eagle_infer_b.py`).
|
||||
- If a single job in a github workflow runs longer than 30 mins, split it into smaller jobs/steps.
|
||||
- Reuse server launches in your unit tests to make tests run faster.
|
||||
- Never use `pickle.loads()`, `pickle.load()`, or `recv_pyobj()` to deserialize untrusted or network-received data. Python's [pickle module is not secure](https://docs.python.org/3/library/pickle.html) — it can execute arbitrary code during deserialization. Use safe serialization formats such as [msgpack](https://github.com/jcrist/msgspec) or JSON instead.
|
||||
- When supporting new hardware or features, follow these guidelines:
|
||||
- Do not drastically change existing code.
|
||||
- Always prefer new files to introduce specific components for your new hardware (e.g., `allocator_ascend.py`).
|
||||
- If you write multiple if/else blocks for new features, ensure the common path (e.g., NVIDIA hardware or the existing code path) is the first branch.
|
||||
|
||||
## How to update sgl-kernel
|
||||
Since sglang and sgl-kernel are separate Python packages, our current GitHub CI infrastructure does not support updating a kernel and using it immediately within the same pull request (PR).
|
||||
To add a new kernel or modify an existing one in the sgl-kernel package, you must use multiple PRs.
|
||||
Since sglang and the `sglang-kernel` (prior `sgl-kernel`) distribution are separate Python packages, our current GitHub CI infrastructure does not support updating a kernel and using it immediately within the same pull request (PR).
|
||||
To add a new kernel or modify an existing one in the `sgl-kernel/` source tree, you must use multiple PRs.
|
||||
|
||||
Follow these steps:
|
||||
|
||||
1. Submit a PR to update the sgl-kernel source code without using it in sglang python package (e.g., [#8884](https://github.com/sgl-project/sglang/pull/8884/files)).
|
||||
2. Bump the version of sgl-kernel (e.g., [#9220](https://github.com/sgl-project/sglang/pull/9220/files)).
|
||||
- Once merged, this will trigger an automatic release of the sgl-kernel wheel to PyPI.
|
||||
2. Bump the version of the kernel package (e.g., [#9220](https://github.com/sgl-project/sglang/pull/9220/files)).
|
||||
- Once merged, this will trigger an automatic release of the `sglang-kernel` wheel to PyPI.
|
||||
- If not urgent, you can wait for other people to release the wheel. A new version will typically be released within one week.
|
||||
3. Apply the changes:
|
||||
- Update the sgl-kernel version in `sglang/python/pyproject.toml` to use the modified kernels.
|
||||
- Update the `sglang-kernel` version in `sglang/python/pyproject.toml` to use the modified kernels.
|
||||
- Update the related caller code in the sglang to use the new kernel.
|
||||
|
||||
## Tips for newcomers
|
||||
|
||||
If you want to contribute but don't have a specific idea in mind, pick issues labeled ["good first issue" or "help wanted"](https://github.com/sgl-project/sglang/issues?q=is%3Aissue+label%3A%22good+first+issue%22%2C%22help+wanted%22). These tasks typically have lower complexity and provide an excellent introduction to the codebase. Also check out this [code walk-through](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/sglang/code-walk-through) for a deeper look into SGLang's workflow.
|
||||
If you want to contribute but don’t have a specific idea in mind, pick issues labeled [“good first issue” or “help wanted”](https://github.com/sgl-project/sglang/issues?q=is%3Aissue+label%3A%22good+first+issue%22%2C%22help+wanted%22). These tasks typically have lower complexity and provide an excellent introduction to the codebase.
|
||||
|
||||
Also check out the following materials as startup guide:
|
||||
- [Mini-SGLang](https://github.com/sgl-project/mini-sglang) for a quick overview on the structure of sglang.
|
||||
- [Code Walk-through](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/sglang/code-walk-through) for a deeper look into SGLang’s workflow.
|
||||
- [GTC-2026 Training Lab](https://drive.google.com/file/d/1mwOZEtipNLJzrflCTodj34KhuOZEoEw5/view?usp=drive_link) for hands-on practices of how to do optimization, benchmarking, or profiling on a launched SGLang instance.
|
||||
|
||||
If you have any questions or want to start a discussion, please feel free to ask in our [Slack channel](https://slack.sglang.io).
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ metatags:
|
||||
## Setup VSCode on a Remote Host
|
||||
(Optional - you can skip this step if you plan to run sglang dev container locally)
|
||||
|
||||
1. In the remote host, download `code` from [Https://code.visualstudio.com/docs/?dv=linux64cli](https://code.visualstudio.com/download) and run `code tunnel` in a shell.
|
||||
1. In the remote host, download `code` from [VSCode](https://code.visualstudio.com/download) and run `code tunnel` in a shell.
|
||||
|
||||
Example
|
||||
```bash Command
|
||||
|
||||
+425
-266
@@ -1,266 +1,425 @@
|
||||
---
|
||||
title: "Development Guide for JIT Kernels"
|
||||
sidebarTitle: "JIT Kernels"
|
||||
metatags:
|
||||
description: "SGLang JIT kernel development: clangd setup, TensorMatcher, LaunchKernel, add_constant example walkthrough."
|
||||
---
|
||||
## Environment Setup
|
||||
|
||||
We strongly recommend using `clangd` as the language server for JIT kernel development.
|
||||
For Ubuntu/Debian, you can download clangd from [apt.llvm.org](https://apt.llvm.org/).
|
||||
If you are using VS Code, we recommend installing the `clangd` extension for better IDE integration.
|
||||
|
||||
All JIT-related files are located in `python/sglang/jit_kernel`.
|
||||
Unlike `sgl-kernel`, which compiles CUDA/C++ binaries ahead of time (AOT), just-in-time (JIT) kernels are compiled at runtime.
|
||||
Consequently, a static `compile_commands.json` cannot be generated.
|
||||
To enable code completion with `clangd`, run `python -m sglang.jit_kernel` to generate a `.clangd` configuration file in your current directory.
|
||||
After generating the file, restart the clangd language server. It should now recognize all JIT kernel files.
|
||||
|
||||
## Code Structure
|
||||
|
||||
### C++ Implementation
|
||||
|
||||
C++ source code is located in `python/sglang/jit_kernel/csrc`.
|
||||
Reusable functions should be placed in `python/sglang/jit_kernel/include`.
|
||||
|
||||
We use [tvm-ffi](https://github.com/apache/tvm-ffi) for efficient foreign language bindings.
|
||||
Refer to the [documentation](https://tvm.apache.org/ffi/) for advanced usage, such as exporting C++ objects.
|
||||
Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from Python.
|
||||
|
||||
### Python Interface
|
||||
|
||||
Python interfaces are defined in `python/sglang/jit_kernel`.
|
||||
The `load_jit` utility function in `python/sglang/jit_kernel/utils.py` loads and returns the compiled module.
|
||||
To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`.
|
||||
The function can then be called in Python as `module.func`.
|
||||
|
||||
### C++ Utilities
|
||||
|
||||
The following C++ utilities are available:
|
||||
|
||||
#### Integer Range
|
||||
|
||||
Similar to PyTorch, we provide an `irange` function to represent an integer range.
|
||||
|
||||
```C++ Example
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
void test() {
|
||||
for (auto i : host::irange(100)) { // [0, 100)
|
||||
// do something
|
||||
}
|
||||
for (auto i : host::irange(0, 100)) { // [0, 100)
|
||||
// do something
|
||||
}
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
#### Runtime Checking
|
||||
|
||||
`RuntimeCheck` validates conditions at runtime. It accepts optional arguments for error reporting.
|
||||
If the check fails, these arguments are output to aid debugging.
|
||||
`RuntimeDeviceCheck` verifies the status of the last kernel launch.
|
||||
|
||||
```C++ Example
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
void test() {
|
||||
host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2);
|
||||
host::RuntimeDeviceCheck();
|
||||
// check the provided `cudaError_t`
|
||||
host::RuntimeDeviceCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
#### Tensor Checking
|
||||
|
||||
`TensorMatcher` provides a readable way to validate and extract tensor shape information.
|
||||
|
||||
```cpp Example
|
||||
#include <sgl_kernel/tensor.h>
|
||||
|
||||
void test(const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache) {
|
||||
using namespace host;
|
||||
|
||||
auto D = SymbolicSize{"D"}; // cache dimension
|
||||
auto N = SymbolicSize{"N"}; // kvcache stride
|
||||
auto dtype = SymbolicDType{};
|
||||
auto device = SymbolicDevice{};
|
||||
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_strides({N, 1})
|
||||
.with_dtype<int32_t, int64_t>(dtype)
|
||||
.with_device<kDLCUDA, kDLCPU>(device)
|
||||
.verify(k_cache)
|
||||
.verify(v_cache);
|
||||
}
|
||||
```
|
||||
|
||||
Configure the `TensorMatcher` with expected stride, dtype, and device properties before verification.
|
||||
- If `with_strides` is omitted, the tensor is expected to be contiguous.
|
||||
- Template arguments in `with_dtype` restrict the allowed data types.
|
||||
- Template arguments in `with_device` restrict the allowed devices.
|
||||
- Values passed to `with_xxx` methods enforce equality checks.
|
||||
- Passing `-1` for size or stride allows matching any value.
|
||||
|
||||
A `Symbolic` variable must resolve to the same value across all verifications.
|
||||
Use `.unwrap()` to retrieve the matched value after verification.
|
||||
|
||||
<Note>
|
||||
`TensorMatcher` is a temporary expression and should not be stored in a variable.
|
||||
</Note>
|
||||
|
||||
<Tip>
|
||||
Add `//` at the end of the `TensorMatcher` chain to enforce proper indentation.
|
||||
</Tip>
|
||||
|
||||
#### Kernel Launching
|
||||
|
||||
`LaunchKernel::resolve_device` retrieves the current `cudaStream` from PyTorch.
|
||||
Kernels can also be launched directly using `LaunchKernel`.
|
||||
|
||||
```cpp Example
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
|
||||
__global__ void kernel() {}
|
||||
|
||||
void test() {
|
||||
const auto num_blocks = 1;
|
||||
const auto num_threads = 32;
|
||||
const auto dynamic_smem = 0;
|
||||
|
||||
DLDevice dev; // suppose this is initialized properly
|
||||
host::LaunchKernel(num_blocks, num_threads, dev)(kernel);
|
||||
|
||||
cudaStream_t stream = host::LaunchKernel::resolve_device(dev);
|
||||
host::LaunchKernel(num_blocks, num_threads, stream, dynamic_smem)(kernel);
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
## Add new kernels
|
||||
|
||||
This section walks through a complete, end-to-end example of adding a new JIT kernel to the system.
|
||||
We use a simple add_constant kernel as a running example, which adds a constant integer value to every element of an input tensor.
|
||||
|
||||
Conceptually, the Python interface looks like this:
|
||||
|
||||
```python Example
|
||||
def add_constant(src: torch.Tensor, c: int):
|
||||
return src + c
|
||||
```
|
||||
|
||||
### STEP 1: Write the C++ kernel
|
||||
|
||||
Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
|
||||
|
||||
```cpp Example
|
||||
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel
|
||||
#include <sgl_kernel/utils.h> // For div_ceil, RuntimeCheck
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
template <int32_t kConstant>
|
||||
__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
|
||||
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < length) {
|
||||
dst[idx] = src[idx] + kConstant;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr size_t kBlockSize = 256;
|
||||
|
||||
// You can also use struct with static method as an alternative
|
||||
template <int32_t kConstant>
|
||||
void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
|
||||
using namespace host;
|
||||
|
||||
// 1. Validate input tensors
|
||||
SymbolicSize N = {"num_elements"};
|
||||
SymbolicDevice device_;
|
||||
TensorMatcher({N}) // 1D tensor, must be contiguous
|
||||
.with_dtype<int32_t>() // must be int32
|
||||
.with_device<kDLCUDA>(device_) // must be on CUDA device
|
||||
.verify(dst) // check tensor dst
|
||||
.verify(src); // check tensor src
|
||||
|
||||
// 2. Extract required parameters, prepare for kernel launch
|
||||
const size_t num_elements = N.unwrap();
|
||||
const size_t grid_size = div_ceil(num_elements, kBlockSize);
|
||||
const DLDevice device = device_.unwrap();
|
||||
// some extra runtime checks using host::RuntimeCheck
|
||||
RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements);
|
||||
|
||||
// 3. Launch the kernel. Error code will be automatically checked.
|
||||
LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
|
||||
// kernel function
|
||||
add_constant_kernel<kConstant>,
|
||||
// kernel arguments
|
||||
static_cast<int32_t*>(dst.data_ptr()),
|
||||
static_cast<int32_t*>(src.data_ptr()),
|
||||
num_elements);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
```
|
||||
|
||||
### STEP 2: Create Python Interfaces
|
||||
|
||||
Next, expose the kernel through a Python wrapper.
|
||||
Create a new file at [jit_kernel/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/add_constant.py) and expose the needed interfaces.
|
||||
|
||||
```python Example
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _jit_add_constant_module(constant: int) -> Module:
|
||||
args = make_cpp_args(constant) # pass all the template argument
|
||||
return load_jit(
|
||||
"add_constant",
|
||||
*args,
|
||||
cuda_files=["add_constant.cuh"],
|
||||
cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
|
||||
)
|
||||
|
||||
|
||||
def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
|
||||
dst = torch.empty_like(src)
|
||||
module = _jit_add_constant_module(constant)
|
||||
module.add_constant(dst, src)
|
||||
return dst
|
||||
|
||||
```
|
||||
|
||||
### STEP 3: Use your kernel
|
||||
|
||||
Finally, import and use the kernel like a regular Python function:
|
||||
|
||||
```python Example
|
||||
from sglang.jit_kernel.add_constant import add_constant
|
||||
```
|
||||
|
||||
For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/test_add_constant.py).
|
||||
---
|
||||
title: "Development Guide for JIT Kernels"
|
||||
sidebarTitle: "JIT Kernels"
|
||||
metatags:
|
||||
description: "SGLang JIT kernel development: clangd setup, TensorMatcher, LaunchKernel, add_constant example walkthrough."
|
||||
---
|
||||
## Environment Setup
|
||||
|
||||
We strongly recommend using `clangd` as the language server for JIT kernel development.
|
||||
For Ubuntu/Debian, you can download clangd from [apt.llvm.org](https://apt.llvm.org/).
|
||||
If you are using VS Code, we recommend installing the `clangd` extension for better IDE integration.
|
||||
|
||||
All JIT-related files are located in `python/sglang/jit_kernel`.
|
||||
Unlike `sgl-kernel`, which compiles CUDA/C++ binaries ahead of time (AOT), just-in-time (JIT) kernels are compiled at runtime.
|
||||
Consequently, a static `compile_commands.json` cannot be generated.
|
||||
To enable code completion with `clangd`, run `python -m sglang.jit_kernel` to generate a `.clangd` configuration file in your current directory.
|
||||
After generating the file, restart the clangd language server. It should now recognize all JIT kernel files.
|
||||
|
||||
## Code Structure
|
||||
|
||||
### C++ Implementation
|
||||
|
||||
C++ source code is located in `python/sglang/jit_kernel/csrc`.
|
||||
Reusable functions should be placed in `python/sglang/jit_kernel/include`.
|
||||
|
||||
We use [tvm-ffi](https://github.com/apache/tvm-ffi) for efficient foreign language bindings.
|
||||
Refer to the [documentation](https://tvm.apache.org/ffi/) for advanced usage, such as exporting C++ objects.
|
||||
Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from Python.
|
||||
|
||||
### Python Interface
|
||||
|
||||
Python interfaces are defined in `python/sglang/jit_kernel`.
|
||||
The `load_jit` utility function in `python/sglang/jit_kernel/utils.py` loads and returns the compiled module.
|
||||
To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`.
|
||||
The function can then be called in Python as `module.func`.
|
||||
|
||||
For caching compiled modules, prefer `sglang.jit_kernel.utils.cache_once` over `functools.lru_cache`.
|
||||
`functools.lru_cache` is not compatible with `torch.compile`.
|
||||
|
||||
### C++ Utilities
|
||||
|
||||
The following C++ utilities are available:
|
||||
|
||||
#### Integer Range
|
||||
|
||||
Similar to PyTorch, we provide an `irange` function to represent an integer range.
|
||||
|
||||
```C++ Example
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
void test() {
|
||||
for (auto i : host::irange(100)) { // [0, 100)
|
||||
// do something
|
||||
}
|
||||
for (auto i : host::irange(0, 100)) { // [0, 100)
|
||||
// do something
|
||||
}
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
#### Runtime Checking
|
||||
|
||||
`RuntimeCheck` validates conditions at runtime. It accepts optional arguments for error reporting.
|
||||
If the check fails, these arguments are output to aid debugging.
|
||||
`RuntimeDeviceCheck` verifies the status of the last kernel launch.
|
||||
|
||||
```C++ Example
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
void test() {
|
||||
host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2);
|
||||
host::RuntimeDeviceCheck();
|
||||
// check the provided `cudaError_t`
|
||||
host::RuntimeDeviceCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
#### Tensor Checking
|
||||
|
||||
`TensorMatcher` provides a readable way to validate and extract tensor shape information.
|
||||
|
||||
```cpp Example
|
||||
#include <sgl_kernel/tensor.h>
|
||||
|
||||
void test(const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache) {
|
||||
using namespace host;
|
||||
|
||||
auto D = SymbolicSize{"D"}; // cache dimension
|
||||
auto N = SymbolicSize{"N"}; // kvcache stride
|
||||
auto dtype = SymbolicDType{};
|
||||
auto device = SymbolicDevice{};
|
||||
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_strides({N, 1})
|
||||
.with_dtype<int32_t, int64_t>(dtype)
|
||||
.with_device<kDLCUDA, kDLCPU>(device)
|
||||
.verify(k_cache)
|
||||
.verify(v_cache);
|
||||
}
|
||||
```
|
||||
|
||||
Configure the `TensorMatcher` with expected stride, dtype, and device properties before verification.
|
||||
- If `with_strides` is omitted, the tensor is expected to be contiguous.
|
||||
- Template arguments in `with_dtype` restrict the allowed data types.
|
||||
- Template arguments in `with_device` restrict the allowed devices.
|
||||
- Values passed to `with_xxx` methods enforce equality checks.
|
||||
- Passing `-1` for size or stride allows matching any value.
|
||||
|
||||
A `Symbolic` variable must resolve to the same value across all verifications.
|
||||
Use `.unwrap()` to retrieve the matched value after verification.
|
||||
|
||||
> Note: `TensorMatcher` is a temporary expression and should not be stored in a variable.
|
||||
|
||||
> Tip: Add `//` at the end of the `TensorMatcher` chain to enforce proper indentation.
|
||||
|
||||
#### Kernel Launching
|
||||
|
||||
`LaunchKernel::resolve_device` retrieves the current `cudaStream` from PyTorch.
|
||||
Kernels can also be launched directly using `LaunchKernel`.
|
||||
|
||||
```cpp Example
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
|
||||
__global__ void kernel() {}
|
||||
|
||||
void test() {
|
||||
const auto num_blocks = 1;
|
||||
const auto num_threads = 32;
|
||||
const auto dynamic_smem = 0;
|
||||
|
||||
DLDevice dev; // suppose this is initialized properly
|
||||
host::LaunchKernel(num_blocks, num_threads, dev)(kernel);
|
||||
|
||||
cudaStream_t stream = host::LaunchKernel::resolve_device(dev);
|
||||
host::LaunchKernel(num_blocks, num_threads, stream, dynamic_smem)(kernel);
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
## Add new kernels
|
||||
|
||||
This section walks through a complete, end-to-end example of adding a new JIT kernel to the system.
|
||||
We use a simple add_constant kernel as a running example, which adds a constant integer value to every element of an input tensor.
|
||||
|
||||
Conceptually, the Python interface looks like this:
|
||||
|
||||
```python Example
|
||||
def add_constant(src: torch.Tensor, c: int):
|
||||
return src + c
|
||||
```
|
||||
|
||||
### STEP 1: Write the C++ kernel
|
||||
|
||||
Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
|
||||
|
||||
```cpp Example
|
||||
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel
|
||||
#include <sgl_kernel/utils.h> // For div_ceil, RuntimeCheck
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
template <int32_t kConstant>
|
||||
__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
|
||||
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < length) {
|
||||
dst[idx] = src[idx] + kConstant;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr size_t kBlockSize = 256;
|
||||
|
||||
// You can also use struct with static method as an alternative
|
||||
template <int32_t kConstant>
|
||||
void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
|
||||
using namespace host;
|
||||
|
||||
// 1. Validate input tensors
|
||||
SymbolicSize N = {"num_elements"};
|
||||
SymbolicDevice device_;
|
||||
TensorMatcher({N}) // 1D tensor, must be contiguous
|
||||
.with_dtype<int32_t>() // must be int32
|
||||
.with_device<kDLCUDA>(device_) // must be on CUDA device
|
||||
.verify(dst) // check tensor dst
|
||||
.verify(src); // check tensor src
|
||||
|
||||
// 2. Extract required parameters, prepare for kernel launch
|
||||
const size_t num_elements = N.unwrap();
|
||||
const size_t grid_size = div_ceil(num_elements, kBlockSize);
|
||||
const DLDevice device = device_.unwrap();
|
||||
// some extra runtime checks using host::RuntimeCheck
|
||||
RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements);
|
||||
|
||||
// 3. Launch the kernel. Error code will be automatically checked.
|
||||
LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
|
||||
// kernel function
|
||||
add_constant_kernel<kConstant>,
|
||||
// kernel arguments
|
||||
static_cast<int32_t*>(dst.data_ptr()),
|
||||
static_cast<int32_t*>(src.data_ptr()),
|
||||
num_elements);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
```
|
||||
|
||||
### STEP 2: Create Python Interfaces
|
||||
|
||||
Next, expose the kernel through a Python wrapper.
|
||||
Create a new file at [jit_kernel/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/add_constant.py) and expose the needed interfaces.
|
||||
|
||||
```python Example
|
||||
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_add_constant_module(constant: int) -> Module:
|
||||
args = make_cpp_args(constant) # pass all the template argument
|
||||
return load_jit(
|
||||
"add_constant",
|
||||
*args,
|
||||
cuda_files=["add_constant.cuh"],
|
||||
cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
|
||||
)
|
||||
|
||||
|
||||
def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
|
||||
if not src.is_cuda:
|
||||
raise RuntimeError("src must be a CUDA tensor")
|
||||
if src.dtype != torch.int32:
|
||||
raise RuntimeError(f"Unsupported dtype {src.dtype}. Supported: int32")
|
||||
dst = torch.empty_like(src)
|
||||
module = _jit_add_constant_module(constant)
|
||||
module.add_constant(dst, src)
|
||||
return dst
|
||||
|
||||
```
|
||||
|
||||
Keep the Python wrapper thin, but still validate the basic invariants such as device and dtype before dispatch. In the current JIT/FFI path, invalid tensors are not always rejected safely before launch.
|
||||
|
||||
### STEP 3: Use your kernel
|
||||
|
||||
Finally, import and use the kernel like a regular Python function:
|
||||
|
||||
```python Example
|
||||
from sglang.jit_kernel.add_constant import add_constant
|
||||
```
|
||||
|
||||
For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/tests/test_add_constant.py).
|
||||
|
||||
## C++ Include Library Reference
|
||||
|
||||
The JIT kernel framework provides a set of reusable C++ headers in
|
||||
`python/sglang/jit_kernel/include/sgl_kernel/`. Each header is designed
|
||||
to be lightweight and self-contained. Below is a summary of each header
|
||||
and its key APIs.
|
||||
|
||||
### Core Utilities
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Header</th>
|
||||
<th>Namespace</th>
|
||||
<th>Purpose</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td><code>utils.h</code></td>
|
||||
<td><code>host</code></td>
|
||||
<td>Host-side essentials: <code>RuntimeCheck</code>, <code>Panic</code>, <code>div_ceil</code>, <code>irange</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>utils.cuh</code></td>
|
||||
<td><code>device</code> / <code>host</code></td>
|
||||
<td>Type aliases (<code>fp16_t</code>, <code>bf16_t</code>, ...), <code>SGL_DEVICE</code> macro, PDL helpers, <code>LaunchKernel</code>, <code>RuntimeDeviceCheck</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>source_location.h</code></td>
|
||||
<td>(global)</td>
|
||||
<td>Portable <code>std::source_location</code> wrapper for error reporting</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>runtime.cuh</code></td>
|
||||
<td><code>host::runtime</code></td>
|
||||
<td>CUDA runtime queries: <code>get_blocks_per_sm</code>, <code>get_sm_count</code>, <code>get_cc_major</code>, <code>get_runtime_version</code>, <code>get_available_dynamic_smem_per_block</code></td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
### Tensor Validation
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Header</th>
|
||||
<th>Namespace</th>
|
||||
<th>Purpose</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td><code>tensor.h</code></td>
|
||||
<td><code>host</code></td>
|
||||
<td><code>TensorMatcher</code>, <code>SymbolicSize</code>, <code>SymbolicDType</code>, <code>SymbolicDevice</code></td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
### Math & Type System
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Header</th>
|
||||
<th>Namespace</th>
|
||||
<th>Purpose</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td><code>math.cuh</code></td>
|
||||
<td><code>device::math</code></td>
|
||||
<td><code>max</code>, <code>min</code>, <code>abs</code>, <code>sqrt</code>, <code>rsqrt</code>, <code>exp</code>, <code>sin</code>, <code>cos</code>, constants</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>type.cuh</code></td>
|
||||
<td>(global) / <code>device</code></td>
|
||||
<td><code>dtype_trait<T></code>, <code>packed_t<T></code>, <code>device::cast<To>(from)</code></td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
### Memory Access
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Header</th>
|
||||
<th>Namespace</th>
|
||||
<th>Purpose</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td><code>vec.cuh</code></td>
|
||||
<td><code>device</code></td>
|
||||
<td><code>AlignedVector<T, N></code> - vectorized load/store (up to 128-bit; 256-bit requires Blackwell GPUs)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>tile.cuh</code></td>
|
||||
<td><code>device::tile</code></td>
|
||||
<td><code>Memory<T></code> - cooperative tiled memory I/O (thread/warp/CTA)</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
### Parallel Primitives
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Header</th>
|
||||
<th>Namespace</th>
|
||||
<th>Purpose</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td><code>warp.cuh</code></td>
|
||||
<td><code>device::warp</code></td>
|
||||
<td><code>reduce_sum</code>, <code>reduce_max</code> via <code>__shfl_xor_sync</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>cta.cuh</code></td>
|
||||
<td><code>device::cta</code></td>
|
||||
<td><code>reduce_max</code> across warps via shared memory</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>atomic.cuh</code></td>
|
||||
<td><code>device::atomic</code></td>
|
||||
<td><code>max</code> - atomic float max (CUDA + ROCm fallback)</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
### Reusable Kernel Templates
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Header</th>
|
||||
<th>Namespace</th>
|
||||
<th>Purpose</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td><code>impl/norm.cuh</code></td>
|
||||
<td><code>host::norm</code> / <code>device::norm</code></td>
|
||||
<td>RMSNorm building blocks (warp & CTA paths, <code>StorageType</code>)</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
@@ -5,7 +5,7 @@ description: Contributing to SGLang — development setup, benchmarking, and eva
|
||||
|
||||
- [Contribution Guide](./contribution_guide)
|
||||
- [Development Guide (Docker)](./development_guide_using_docker)
|
||||
- [JIT Kernels](./JIT_kernels)
|
||||
- [JIT Kernels](./development_jit_kernel_guide)
|
||||
- [Benchmark and Profiling](./benchmark_and_profiling)
|
||||
- [Bench Serving](./bench_serving)
|
||||
- [Evaluating New Models](./evaluating_new_models)
|
||||
|
||||
@@ -10,14 +10,14 @@ metatags:
|
||||
**You can mount a folder for the shared huggingface model weights cache. **
|
||||
The command below uses `/tmp/huggingface` as an example.
|
||||
|
||||
```text Output
|
||||
```
|
||||
docker pull nvidia/cuda:12.9.1-devel-ubuntu22.04
|
||||
# Nvidia
|
||||
docker run --shm-size 128g -it -v /tmp/huggingface:/hf_home --gpus all nvidia/cuda:12.9.1-devel-ubuntu22.04 /bin/bash
|
||||
# AMD
|
||||
docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.0rc1-rocm630 /bin/bash
|
||||
docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.8-rocm700-mi30x /bin/bash
|
||||
# AMD just the last 2 GPUs
|
||||
docker run --rm --device=/dev/kfd --device=/dev/dri/renderD176 --device=/dev/dri/renderD184 --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.0rc1-rocm630 /bin/bash
|
||||
docker run --rm --device=/dev/kfd --device=/dev/dri/renderD176 --device=/dev/dri/renderD184 --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.8-rocm700-mi30x /bin/bash
|
||||
```
|
||||
|
||||
### Step 2: Configure the runner by `config.sh`
|
||||
@@ -30,11 +30,11 @@ pip install --upgrade pip
|
||||
export RUNNER_ALLOW_RUNASROOT=1
|
||||
```
|
||||
|
||||
Then follow https://github.com/sgl-project/sglang/settings/actions/runners/new?arch=x64&os=linux to run `config.sh`
|
||||
Then follow https://docs.github.com/en/actions/hosting-your-own-runners/managing-self-hosted-runners/adding-self-hosted-runners to run `config.sh`
|
||||
|
||||
**Notes**
|
||||
- Do not need to specify the runner group
|
||||
- Give it a name (e.g., `test-sgl-gpu-0`) and some labels (e.g., `1-gpu-runner`). The labels can be edited later in Github Settings.
|
||||
- Give it a name (e.g., `test-sgl-gpu-0`) and some labels (e.g., `1-gpu-h100`). The labels can be edited later in Github Settings.
|
||||
- Do not need to change the work folder.
|
||||
|
||||
### Step 3: Run the runner by `run.sh`
|
||||
|
||||
Reference in New Issue
Block a user