chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add SGLang CUDA crash API logging inspired by FlashInfer (#20910)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-22 16:39:40 +08:00
committed by GitHub
parent bb737d7a82
commit 766d225fcc
46 changed files with 1585 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/debug-cuda-crash/SKILL.md
+657
View File
@@ -0,0 +1,657 @@
---
name: debug-cuda-crash
description: Call this skill when you need to debug CUDA crashes in SGLang using kernel API logging
---
# Tutorial: Debugging CUDA Crashes with Kernel API Logging
This tutorial shows you how to debug CUDA crashes and errors in SGLang using the `@debug_kernel_api` logging decorator.
## Goal
When your code crashes with CUDA errors such as illegal memory access, device-side assert, out-of-bounds, or NaN/Inf, use kernel API logging to:
- Capture input tensors BEFORE the crash occurs
- Understand what data caused the problem
- Track tensor shapes, dtypes, and values through the call boundary that triggered the crash
- Detect numerical issues such as NaN, Inf, or obviously wrong shapes
## Why Use Kernel API Logging?
**Problem**: CUDA errors often crash the program before normal debugging output is flushed.
**Solution**: SGLang's `@debug_kernel_api` decorator logs inputs before execution, so you can still see what caused the crash even after the program aborts.
## What Is Covered?
The current logging coverage focuses on the highest-value kernel boundaries in SGLang:
- Custom ops registered through `register_custom_op(...)`
- External custom ops registered through `register_custom_op_from_extern(...)`
- LLM attention, linear, quantization, and multi-platform wrapper entry points
- Diffusion attention impl, linear, rotary, and custom-op wrapper entry points
- Selected direct `torch.ops.sglang.*` hotspots and model-specific bypasses
This means the logging is useful for both LLM and diffusion kernel debugging, but it does not automatically cover every pure PyTorch call in the repository.
## Step 1: Enable Kernel API Logging
### Basic Logging (Function Names Only)
```bash
export SGLANG_KERNEL_API_LOGLEVEL=1
export SGLANG_KERNEL_API_LOGDEST=stdout
python my_script.py
```
Output:
```
================================================================================
[2026-03-19 00:47:06] SGLang Kernel API Call: RMSNorm.forward
================================================================================
[2026-03-19 00:47:06] SGLang Kernel API Call: sglang.quant_method.UnquantizedLinearMethod.apply
================================================================================
[2026-03-19 00:47:06] SGLang Kernel API Call: sglang.custom_op.fused_inplace_qknorm
```
This is a real level-1 excerpt captured from `Qwen/Qwen3-0.6B`.
### Detailed Logging (Inputs with Metadata)
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
export SGLANG_KERNEL_API_LOGDEST=debug.log
python my_script.py
```
Output in `debug.log`:
```
================================================================================
[2026-03-19 00:47:30] SGLang Kernel API Call: sglang.quant_method.UnquantizedLinearMethod.apply
Positional input arguments:
arg[0]=QKVParallelLinear(
repr=QKVParallelLinear(in_features=1024, output_features=4096, bias=False, tp_size=1, gather_output=False)
)
arg[1]=Tensor(
shape=(1, 1024)
dtype=torch.bfloat16
device=cuda:0
requires_grad=False
is_contiguous=True
)
arg[2]=None
Output:
return=Tensor(
shape=(1, 4096)
dtype=torch.bfloat16
device=cuda:0
requires_grad=False
is_contiguous=True
)
```
This is a real level-3 excerpt captured from `Qwen/Qwen3-0.6B`.
### Full Logging (With Tensor Statistics)
```bash
export SGLANG_KERNEL_API_LOGLEVEL=5
export SGLANG_KERNEL_API_LOGDEST=debug.log
python my_script.py
```
Additional output:
```
================================================================================
[2026-03-19 01:00:42] SGLang Kernel API Call: diffusion.quant_method.UnquantizedLinearMethod.apply
Positional input arguments:
arg[1]=Tensor(
shape=(1, 77, 768)
dtype=torch.bfloat16
device=cuda:0
requires_grad=False
is_contiguous=True
min=-27.250000
max=28.500000
mean=0.011723
nan_count=0
inf_count=0
)
Output:
return=Tensor(
shape=(1, 77, 2304)
dtype=torch.bfloat16
device=cuda:0
requires_grad=False
is_contiguous=True
min=-8.937500
max=9.375000
mean=0.009460
nan_count=0
inf_count=0
)
```
This is a real level-5 excerpt captured from `black-forest-labs/FLUX.1-dev`.
### Crash-Safe Dumps (Inputs Saved Before Execution)
```bash
export SGLANG_KERNEL_API_LOGLEVEL=10
export SGLANG_KERNEL_API_LOGDEST=debug.log
export SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_kernel_api_dumps
python my_script.py
```
At level 10, SGLang saves the inputs before execution. If the kernel crashes, the dump directory still contains the inputs and exception metadata.
If CUDA graph capture is active, tensor dumps are skipped automatically to avoid capture-time CUDA errors. In that case, you still get the kernel API call log, but not `inputs.pt` / `outputs.pt`.
Level-10 dumps are best understood as crash-safe call snapshots. They always preserve the observed call boundary. They do not guarantee one-click replay for every method, because some methods depend on module state that is not serialized into the dump.
Real level-10 dump layout from `Qwen/Qwen3-0.6B`:
```text
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001/inputs.pt
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001/metadata.json
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001/outputs.pt
```
Real `metadata.json` excerpt:
```json
{
"function_name": "RotaryEmbedding.forward",
"timestamp": "20260319_004821_182",
"process_id": 919286,
"execution_status": "completed",
"input_tensor_keys": ["arg_0", "arg_1", "arg_2"],
"output_tensor_keys": ["result_0", "result_1"]
}
```
## Step 2: Reproduce an LLM CUDA Crash
Create a temporary reproducer:
```bash
python3 - <<'PY'
from pathlib import Path
Path("/tmp/sglang_llm_crash.py").write_text(
"import torch\\n"
"import torch.nn.functional as F\\n"
"from sglang.srt.utils.custom_op import register_custom_op\\n\\n"
"def _fake_embedding(indices, table):\\n"
" return torch.empty((*indices.shape, table.shape[-1]), device=table.device, dtype=table.dtype)\\n\\n"
"@register_custom_op(op_name='mock_llm_cuda_crash', fake_impl=_fake_embedding)\\n"
"def mock_llm_cuda_crash(indices, table):\\n"
" out = F.embedding(indices, table)\\n"
" torch.cuda.synchronize()\\n"
" return out\\n\\n"
"table = torch.randn(4, 8, device='cuda', dtype=torch.float16)\\n"
"indices = torch.tensor([0, 7], device='cuda', dtype=torch.long)\\n"
"mock_llm_cuda_crash(indices, table)\\n"
)
PY
SGLANG_KERNEL_API_LOGLEVEL=1 \
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_llm_level1.log \
python3 /tmp/sglang_llm_crash.py
```
What to expect:
- The script exits with a CUDA `device-side assert`
- The log still contains the last API boundary before the crash
Try the same example at level 3:
```bash
SGLANG_KERNEL_API_LOGLEVEL=3 \
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_llm_level3.log \
python3 /tmp/sglang_llm_crash.py
```
Now the log shows tensor metadata before the crash.
Try level 10:
```bash
SGLANG_KERNEL_API_LOGLEVEL=10 \
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_llm_level10.log \
SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_llm_level10_dumps \
python3 /tmp/sglang_llm_crash.py
```
Now you should see:
- A log entry for `sglang.custom_op.mock_llm_cuda_crash`
- A dump directory with `inputs.pt`
- `metadata.json` showing `execution_status: "exception"`
- No `outputs.pt`, because the kernel crashed before producing output
For real-model success-path level-10 dumps, it is often easier to temporarily disable CUDA graph and piecewise CUDA graph for the debug run.
## Step 3: Reproduce a Diffusion CUDA Crash
Create a temporary diffusion-side reproducer:
```bash
python3 - <<'PY'
from pathlib import Path
Path("/tmp/sglang_diffusion_crash.py").write_text(
"import torch\\n"
"import torch.nn.functional as F\\n"
"from sglang.multimodal_gen.runtime.layers.utils import register_custom_op\\n\\n"
"def _fake_embedding(positions, cache):\\n"
" return torch.empty((*positions.shape, cache.shape[-1]), device=cache.device, dtype=cache.dtype)\\n\\n"
"@register_custom_op(op_name='mock_diffusion_cuda_crash', fake_impl=_fake_embedding)\\n"
"def mock_diffusion_cuda_crash(positions, cache):\\n"
" out = F.embedding(positions, cache)\\n"
" torch.cuda.synchronize()\\n"
" return out\\n\\n"
"cache = torch.randn(4, 64, device='cuda', dtype=torch.float16)\\n"
"positions = torch.tensor([0, 9], device='cuda', dtype=torch.long)\\n"
"mock_diffusion_cuda_crash(positions, cache)\\n"
)
PY
SGLANG_KERNEL_API_LOGLEVEL=1 \
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_diffusion_level1.log \
python3 /tmp/sglang_diffusion_crash.py
```
Try level 3:
```bash
SGLANG_KERNEL_API_LOGLEVEL=3 \
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_diffusion_level3.log \
python3 /tmp/sglang_diffusion_crash.py
```
Try level 10:
```bash
SGLANG_KERNEL_API_LOGLEVEL=10 \
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_diffusion_level10.log \
SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_diffusion_level10_dumps \
python3 /tmp/sglang_diffusion_crash.py
```
If your local environment has unrelated FlashInfer import issues, resolve them in the shell before running the example. The example itself does not set any `FLASHINFER_*` environment variable.
## Step 4: Multi-Process Debugging
When running with multiple GPUs or worker processes, use `%i` in the log path:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
export SGLANG_KERNEL_API_LOGDEST=debug_rank_%i.log
torchrun --nproc_per_node=4 my_script.py
```
This creates separate logs such as:
- `debug_rank_12345.log`
- `debug_rank_12346.log`
- `debug_rank_12347.log`
- `debug_rank_12348.log`
Real multi-process example from a 2-GPU `Qwen/Qwen2.5-0.5B-Instruct` run:
```text
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950201.log
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950349.log
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950350.log
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950351.log
```
You should usually do the same for level-10 dump directories:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=10
export SGLANG_KERNEL_API_LOGDEST=debug_rank_%i.log
export SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_kernel_api_dumps_%i
```
This avoids multiple ranks writing into the same dump directory tree.
## Step 5: Filter Level-10 Dumps
If level 10 is too noisy, restrict dumps to specific APIs:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=10
export SGLANG_KERNEL_API_LOGDEST=debug.log
export SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_kernel_api_dumps
export SGLANG_KERNEL_API_DUMP_INCLUDE='sglang.custom_op.*'
export SGLANG_KERNEL_API_DUMP_EXCLUDE='*.fake_impl'
```
`SGLANG_KERNEL_API_DUMP_INCLUDE` and `SGLANG_KERNEL_API_DUMP_EXCLUDE` use shell-style wildcard matching.
## Step 6: Common CUDA Errors and What to Check
### Illegal Memory Access or Device-Side Assert
**Typical errors**:
```
RuntimeError: CUDA error: an illegal memory access was encountered
torch.AcceleratorError: CUDA error: device-side assert triggered
```
Use:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
```
Check in the logs:
- ✅ Tensor shapes
- ✅ Tensor dtypes
- ✅ CUDA vs CPU device placement
- ✅ Tensor stride / contiguity
- ✅ Whether the failing call has inputs logged but no outputs logged
Typical shape-mismatch pattern:
```text
SGLang Kernel API Call: ...
arg[0]=Tensor(shape=(..., 128), ...) # ✅ expected dimension
arg[1]=Tensor(shape=(..., 64), ...) # ❌ mismatch
```
This often points to head-dim, hidden-dim, or cache-layout mismatch rather than a random CUDA failure.
### NaN or Inf
Use:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=5
```
Check:
- `min`
- `max`
- `mean`
- `nan_count`
- `inf_count`
Typical bad pattern:
```text
Tensor(
...
min=-1234567.000000 # ❌ suspiciously large
max=9876543.000000 # ❌ suspiciously large
mean=nan # ❌ bad
nan_count=128 # ❌ found NaNs
inf_count=0 # ✅ no Infs here
)
```
This usually means the bad values were already present before the crashing kernel.
### Out of Memory
Use:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
```
Check:
- Unexpectedly large tensor shapes
- Batch size
- Sequence length
- Frame count or image resolution in diffusion workloads
Also check whether a supposedly per-token or per-frame tensor accidentally became full-sequence or full-image sized.
Typical bad pattern:
```text
Tensor(
shape=(1024, 8192, 128, 128) # ❌ way too large
...
)
```
### Example: Spot a Shape Bug from the Log
Suppose the failing API log looks like this:
```text
[2026-03-19 00:47:30] SGLang Kernel API Call: RotaryEmbedding.forward
Positional input arguments:
arg[0]=Tensor(shape=(1, 8), dtype=torch.int64, ...)
arg[1]=Tensor(shape=(1, 8, 8, 256), dtype=torch.bfloat16, ...) # ✅ query
arg[2]=Tensor(shape=(1, 8, 4, 64), dtype=torch.bfloat16, ...) # ❌ key head_dim mismatch
```
What this tells you:
- ✅ positions look reasonable
- ✅ query looks plausible
- ❌ key last dimension is inconsistent with the expected rotary/head dimension
That usually means the bug is in projection layout, head packing, or cache format rather than in the rotary kernel itself.
## Step 7: Combine with compute-sanitizer
For harder bugs, combine kernel API logging with CUDA memory checking:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
export SGLANG_KERNEL_API_LOGDEST=debug.log
compute-sanitizer --tool memcheck python3 /tmp/sglang_llm_crash.py
```
Use `debug.log` to see the exact inputs that reached the crashing API boundary.
Typical `compute-sanitizer` output:
```text
========= COMPUTE-SANITIZER
========= Invalid __global__ write of size 4 bytes
========= at 0x1234 in SomeKernel
========= by thread (256,0,0) in block (10,0,0)
========= Address 0x... is out of bounds
```
Use the sanitizer output to identify the failing kernel and use `debug.log` to identify the exact tensors that reached the API boundary right before it.
If you need more synchronous host-side error reporting, you can try `CUDA_LAUNCH_BLOCKING=1` as a separate follow-up experiment. It is not part of the default workflow because it changes execution timing and can hide concurrency-related behavior.
## Step 8: Combine with cuda-gdb
For crashes that need a stack trace instead of only memory diagnostics:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
export SGLANG_KERNEL_API_LOGDEST=debug.log
cuda-gdb --args python3 /tmp/sglang_llm_crash.py
```
Inside `cuda-gdb`:
```text
(cuda-gdb) run
(cuda-gdb) where
```
Then correlate the backtrace with `debug.log`.
## Step 9: Kernel-Level Debugging with printf()
When you own the CUDA kernel, `printf()` is still useful for narrowing down bad indices, bad launch geometry, or broken state propagation.
Basic pattern:
```cpp
__global__ void MyKernel(const float* input, float* output, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (threadIdx.x == 0 && blockIdx.x == 0) {
printf("n=%d input0=%f\n", n, input[0]);
}
if (idx < n) {
output[idx] = input[idx] * 2.0f;
}
}
```
After launch, force the output to flush:
```python
my_kernel(...)
torch.cuda.synchronize()
```
For warp-specialized kernels, do not blindly print only on `threadIdx.x == 0`. Pick one representative thread per warp or per specialization group instead.
### Warp-Specialized Kernels: Choosing the Right Print Thread
Problem:
- `threadIdx.x == 0` only prints from the first warp in the block
- for warp-specialized kernels, that often misses the warp or group that is actually wrong
Better pattern:
```cpp
__global__ void WarpSpecializedKernel(...) {
// Example: first lane of each warp
if ((threadIdx.x % 32) == 0) {
printf("warp=%d\n", threadIdx.x / 32);
}
}
```
Or, if the kernel is organized in larger specialization groups, print once per group instead of once per block.
Common mistake:
```cpp
// Only warp 0 prints
if (threadIdx.x == 0) {
printf("warp=%d\n", threadIdx.x / 32);
}
```
### Quick Reference
| Kernel Type | Print Condition | Notes |
|----------|----------|-------------|
| Simple kernel | `threadIdx.x == 0` | One thread per block is usually enough |
| Warp-specialized kernel | one representative lane per warp | e.g. `threadIdx.x % 32 == 0` |
| Group-specialized kernel | one representative lane per group | choose based on the kernel's scheduling layout |
### Other Kernel Debugging Tools
```cpp
assert(value >= 0.0f && "value must be non-negative");
static_assert(BLOCK_SIZE % 32 == 0, "BLOCK_SIZE must be warp aligned");
```
## Environment Variables Reference
| Variable | Values | Description |
|----------|--------|-------------|
| `SGLANG_KERNEL_API_LOGLEVEL` | `0` | No logging (default) |
| | `1` | Function names only |
| | `3` | Inputs and outputs with metadata |
| | `5` | Level 3 plus tensor statistics |
| | `10` | Level 5 plus crash-safe tensor dumps |
| `SGLANG_KERNEL_API_LOGDEST` | `stdout` | Log to stdout |
| | `stderr` | Log to stderr |
| | `<path>` | Log to file |
| | `log_%i.txt` | `%i` expands to process ID |
| `SGLANG_KERNEL_API_DUMP_DIR` | `<path>` | Directory for level-10 dumps |
| `SGLANG_KERNEL_API_DUMP_INCLUDE` | wildcard list | Only dump matching API names |
| `SGLANG_KERNEL_API_DUMP_EXCLUDE` | wildcard list | Skip matching API names |
## Best Practices
### 1. Start with Level 3
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
```
Level 3 is usually enough to catch wrong shapes, wrong dtypes, and wrong devices.
### 2. Use Level 5 for Numerical Issues
```bash
export SGLANG_KERNEL_API_LOGLEVEL=5
```
Use it when you suspect NaN or Inf values.
### 3. Use Level 10 for Crash Reproduction
```bash
export SGLANG_KERNEL_API_LOGLEVEL=10
```
This is the most useful mode when the process crashes before you can inspect live tensors.
If you need successful input/output dumps from a real model run, temporarily disable CUDA graph for that debug session.
When level 10 is too noisy, pair it with `SGLANG_KERNEL_API_DUMP_INCLUDE` / `SGLANG_KERNEL_API_DUMP_EXCLUDE` instead of dumping every covered API.
### 4. Log to File for Crashes
```bash
export SGLANG_KERNEL_API_LOGDEST=crash.log
```
File logs are safer than stdout when the process aborts.
### 5. Disable Logging in Production
```bash
unset SGLANG_KERNEL_API_LOGLEVEL
```
When disabled, the decorator returns the original callable and adds no runtime logging overhead.
## Troubleshooting
### No Logs Appear
Check:
1. `echo $SGLANG_KERNEL_API_LOGLEVEL`
2. `echo $SGLANG_KERNEL_API_LOGDEST`
3. Whether the failing path goes through a covered API boundary
### Too Much Output
Reduce the level:
```bash
export SGLANG_KERNEL_API_LOGLEVEL=3
```
### Statistics Are Skipped During CUDA Graph Capture
If you see:
```text
statistics=[skipped: CUDA graph capture in progress]
```
That is expected. Level-5 statistics are intentionally skipped during CUDA graph capture to avoid synchronization side effects.
### Tensor Dumps Are Skipped During CUDA Graph Capture
If you see:
```text
Tensor dump skipped: CUDA graph capture in progress
```
That is also expected. Level-10 dumps require copying tensors to CPU, which is not allowed during CUDA graph capture.
docs/diffusion/environment_variables.md
+12
View File
@@ -40,3 +40,15 @@ These variables configure S3-compatible cloud storage for automatically uploadin
| `SGLANG_S3_REGION_NAME` | us-east-1 | AWS region name |
| `SGLANG_S3_ACCESS_KEY_ID` | not set | AWS Access Key ID |
| `SGLANG_S3_SECRET_ACCESS_KEY` | not set | AWS Secret Access Key |
## CUDA Crash Debugging
These variables enable kernel API logging and optional input/output dumps around diffusion CUDA kernel call boundaries. They are useful when tracking down CUDA crashes such as illegal memory access, device-side assert, or shape mismatches in custom kernels.
| Environment Variable | Default | Description |
|----------------------|---------|-------------|
| `SGLANG_KERNEL_API_LOGLEVEL` | `0` | Controls crash-debug kernel API logging. `1` logs API names, `3` logs tensor metadata, `5` adds tensor statistics, and `10` also writes dump snapshots. |
| `SGLANG_KERNEL_API_LOGDEST` | `stdout` | Destination for crash-debug kernel API logs. Use `stdout`, `stderr`, or a file path. `%i` is replaced with the process PID. |
| `SGLANG_KERNEL_API_DUMP_DIR` | `sglang_kernel_api_dumps` | Output directory for level-10 kernel API dumps. `%i` is replaced with the process PID. |
| `SGLANG_KERNEL_API_DUMP_INCLUDE` | not set | Comma-separated wildcard patterns for kernel API names to include in level-10 dumps. |
| `SGLANG_KERNEL_API_DUMP_EXCLUDE` | not set | Comma-separated wildcard patterns for kernel API names to exclude from level-10 dumps. |
docs/references/environment_variables.md
+5
View File
@@ -151,6 +151,11 @@ SGLang supports various environment variables that can be used to configure its
| `SGLANG_TEST_RETRACT_NO_PREFILL_BS` | When SGLANG_TEST_RETRACT is enabled, no prefill is performed if the batch size exceeds SGLANG_TEST_RETRACT_NO_PREFILL_BS. | `2 ** 31` |
| `SGLANG_RECORD_STEP_TIME` | Record step time for profiling | `false` |
| `SGLANG_TEST_REQUEST_TIME_STATS` | Test request time statistics | `false` |
| `SGLANG_KERNEL_API_LOGLEVEL` | Controls crash-debug kernel API logging. `0` disables logging, `1` logs API names, `3` logs tensor metadata, `5` adds tensor statistics, and `10` also writes pre-call dump snapshots. | `0` |
| `SGLANG_KERNEL_API_LOGDEST` | Destination for crash-debug kernel API logs. Use `stdout`, `stderr`, or a file path. `%i` is replaced with the process PID. | `stdout` |
| `SGLANG_KERNEL_API_DUMP_DIR` | Output directory for level-10 kernel API input/output dumps. `%i` is replaced with the process PID. | `sglang_kernel_api_dumps` |
| `SGLANG_KERNEL_API_DUMP_INCLUDE` | Comma-separated wildcard patterns for kernel API names to include in level-10 dumps. | Not set |
| `SGLANG_KERNEL_API_DUMP_EXCLUDE` | Comma-separated wildcard patterns for kernel API names to exclude from level-10 dumps. | Not set |
## Profiling & Benchmarking
python/sglang/jit_kernel/awq_marlin_repack.py
+3
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit
if TYPE_CHECKING:
@@ -19,6 +20,7 @@ def _jit_awq_marlin_repack_module() -> Module:
)
@maybe_wrap_jit_kernel_debug
def awq_marlin_repack(
b_q_weight: torch.Tensor,
size_k: int,
@@ -37,6 +39,7 @@ def awq_marlin_repack(
return out
@maybe_wrap_jit_kernel_debug
def awq_marlin_moe_repack(
b_q_weight: torch.Tensor,
perm: torch.Tensor,
python/sglang/jit_kernel/debug_utils.py
+45
View File
@@ -0,0 +1,45 @@
import os
from typing import Any, Callable, TypeVar, cast, overload
F = TypeVar("F", bound=Callable[..., Any])
def _wrap_jit_kernel_debug(func: F, op_name: str | None = None) -> F:
try:
if int(os.environ.get("SGLANG_KERNEL_API_LOGLEVEL", "0")) == 0:
return func
except Exception:
return func
try:
from sglang.kernel_api_logging import debug_kernel_api
except Exception:
return func
if getattr(func, "_debug_kernel_wrapped", False):
return func
wrapped = debug_kernel_api(func, op_name=op_name)
setattr(wrapped, "_debug_kernel_wrapped", True)
return cast(F, wrapped)
@overload
def maybe_wrap_jit_kernel_debug(func: F) -> F: ...
@overload
def maybe_wrap_jit_kernel_debug(func: F, op_name: str) -> F: ...
@overload
def maybe_wrap_jit_kernel_debug(*, op_name: str | None = None) -> Callable[[F], F]: ...
def maybe_wrap_jit_kernel_debug(
func: F | None = None, op_name: str | None = None
) -> F | Callable[[F], F]:
if func is None:
return lambda wrapped_func: _wrap_jit_kernel_debug(wrapped_func, op_name)
return _wrap_jit_kernel_debug(func, op_name)
python/sglang/jit_kernel/diffusion/triton/norm.py
+55 -2
View File
@@ -5,6 +5,8 @@ import triton # type: ignore
import triton.language as tl # type: ignore
from torch import Tensor
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
# RMSNorm-fp32
def maybe_contiguous_lastdim(x):
@@ -450,6 +452,7 @@ class LayerNormFn:
return y
@maybe_wrap_jit_kernel_debug
def layer_norm_fn(
x,
weight,
@@ -537,6 +540,7 @@ def _norm_infer_kernel(
tl.store(Y + cols, y, mask=cols < N)
@maybe_wrap_jit_kernel_debug
def norm_infer(
x: Tensor,
weight: Optional[Tensor],
@@ -579,6 +583,7 @@ def norm_infer(
return out
@maybe_wrap_jit_kernel_debug
def rms_norm_fn(
x,
weight,
@@ -625,5 +630,53 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
if current_platform.is_mps():
from .mps_fallback import norm_infer_native, rms_norm_fn_native
norm_infer = norm_infer_native
rms_norm_fn = rms_norm_fn_native
@maybe_wrap_jit_kernel_debug
def norm_infer(
x: Tensor,
weight: Optional[Tensor],
bias: Optional[Tensor],
eps: float,
is_rms_norm: bool = False,
out: Optional[Tensor] = None,
):
return norm_infer_native(x, weight, bias, eps, is_rms_norm, out)
@maybe_wrap_jit_kernel_debug
def rms_norm_fn(
x,
weight,
bias,
residual=None,
x1=None,
weight1=None,
bias1=None,
eps=1e-6,
dropout_p=0.0,
rowscale=None,
prenorm=False,
residual_in_fp32=False,
zero_centered_weight=False,
return_dropout_mask=False,
out_dtype=None,
out=None,
residual_out=None,
):
return rms_norm_fn_native(
x,
weight,
bias,
residual,
x1,
weight1,
bias1,
eps,
dropout_p,
rowscale,
prenorm,
residual_in_fp32,
zero_centered_weight,
return_dropout_mask,
out_dtype,
out,
residual_out,
)
python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py
+5 -1
View File
@@ -2,6 +2,7 @@ import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.srt.utils.custom_op import register_custom_op
@@ -35,6 +36,7 @@ def _rms_norm_tiled_onepass(
tl.store(y_blk, x * rstd * w, mask=mask)
@maybe_wrap_jit_kernel_debug
@register_custom_op(op_name="triton_one_pass_rms_norm_cuda", out_shape="x")
def _triton_one_pass_rms_norm_cuda(
x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
@@ -72,4 +74,6 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
if current_platform.is_mps():
from .mps_fallback import triton_one_pass_rms_norm_native
triton_one_pass_rms_norm = triton_one_pass_rms_norm_native
@maybe_wrap_jit_kernel_debug
def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
return triton_one_pass_rms_norm_native(x, w, eps)
python/sglang/jit_kernel/diffusion/triton/rotary.py
+19 -2
View File
@@ -2,6 +2,7 @@ import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -64,6 +65,7 @@ def _rotary_embedding_kernel(
tl.store(output_row_ptr + offsets_x2, o2_vals.to(x2_vals.dtype), mask=mask)
@maybe_wrap_jit_kernel_debug
def apply_rotary_embedding(
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
) -> torch.Tensor:
@@ -110,9 +112,24 @@ def apply_rotary_embedding(
if current_platform.is_npu():
from .npu_fallback import apply_rotary_embedding_native
apply_rotary_embedding = apply_rotary_embedding_native
@maybe_wrap_jit_kernel_debug
def apply_rotary_embedding(
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
interleaved: bool = False,
) -> torch.Tensor:
return apply_rotary_embedding_native(x, cos, sin, interleaved)
if current_platform.is_mps():
from .mps_fallback import apply_rotary_embedding_native
apply_rotary_embedding = apply_rotary_embedding_native
@maybe_wrap_jit_kernel_debug
def apply_rotary_embedding(
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
interleaved: bool = False,
) -> torch.Tensor:
return apply_rotary_embedding_native(x, cos, sin, interleaved)
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
+56 -3
View File
@@ -2,6 +2,7 @@ import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -444,6 +445,7 @@ def fuse_scale_shift_gate_select01_kernel_blc_opt(
tl.store(gate_out_ptr + go_off, gate, mask=mask)
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_kernel(
x: torch.Tensor,
scale: torch.Tensor,
@@ -563,6 +565,7 @@ def fuse_scale_shift_kernel(
return output
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_gate_select01_kernel(
x: torch.Tensor,
scale0: torch.Tensor,
@@ -635,6 +638,7 @@ def fuse_scale_shift_gate_select01_kernel(
return output, gate_out
@maybe_wrap_jit_kernel_debug
def fuse_layernorm_scale_shift_gate_select01_kernel(
x: torch.Tensor,
weight: torch.Tensor | None,
@@ -724,6 +728,7 @@ def fuse_layernorm_scale_shift_gate_select01_kernel(
return output, gate_out
@maybe_wrap_jit_kernel_debug
def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
x: torch.Tensor,
residual: torch.Tensor,
@@ -834,7 +839,19 @@ def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
if current_platform.is_npu():
from .npu_fallback import fuse_scale_shift_native
fuse_scale_shift_kernel = fuse_scale_shift_native
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_kernel(
x: torch.Tensor,
scale: torch.Tensor,
shift: torch.Tensor,
scale_constant: float = 1.0,
block_l: int = 128,
block_c: int = 128,
):
return fuse_scale_shift_native(
x, scale, shift, scale_constant, block_l, block_c
)
if current_platform.is_mps():
from .mps_fallback import (
@@ -842,5 +859,41 @@ if current_platform.is_mps():
fuse_scale_shift_kernel_native,
)
fuse_scale_shift_kernel = fuse_scale_shift_kernel_native
fuse_scale_shift_gate_select01_kernel = fuse_scale_shift_gate_select01_kernel_native
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_kernel(
x: torch.Tensor,
scale: torch.Tensor,
shift: torch.Tensor,
scale_constant: float = 1.0,
block_l: int = 128,
block_c: int = 128,
):
return fuse_scale_shift_kernel_native(
x, scale, shift, scale_constant, block_l, block_c
)
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_gate_select01_kernel(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
block_l: int = 128,
block_c: int = 128,
):
return fuse_scale_shift_gate_select01_kernel_native(
x,
scale0,
shift0,
gate0,
scale1,
shift1,
gate1,
index,
block_l,
block_c,
)
python/sglang/jit_kernel/flash_attention_v4.py
+4
View File
@@ -4,6 +4,8 @@ from typing import Callable, Optional, Tuple, Union
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
try:
from flash_attn.cute import flash_attn_varlen_func as _flash_attn_varlen_func
except Exception as _e: # pragma: no cover
@@ -17,6 +19,7 @@ def _maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
@maybe_wrap_jit_kernel_debug
def flash_attn_varlen_func(
q: torch.Tensor,
k: torch.Tensor,
@@ -89,6 +92,7 @@ def flash_attn_varlen_func(
return result
@maybe_wrap_jit_kernel_debug
def flash_attn_with_kvcache(
q: torch.Tensor,
k_cache: torch.Tensor,
python/sglang/jit_kernel/fused_store_index_cache.py
+2
View File
@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
@@ -64,6 +65,7 @@ def can_use_nsa_fused_store(
return False
@maybe_wrap_jit_kernel_debug
def fused_store_index_k_cache(
key: torch.Tensor,
index_k_with_scale: torch.Tensor,
python/sglang/jit_kernel/gptq_marlin.py
+2
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
@@ -31,6 +32,7 @@ def _or_empty(
return t if t is not None else torch.empty(0, device=device, dtype=dtype)
@maybe_wrap_jit_kernel_debug
def gptq_marlin_gemm(
a: torch.Tensor,
c: Optional[torch.Tensor],
python/sglang/jit_kernel/gptq_marlin_repack.py
+2
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit
if TYPE_CHECKING:
@@ -22,6 +23,7 @@ def _jit_gptq_marlin_repack_module() -> Module:
)
@maybe_wrap_jit_kernel_debug
def gptq_marlin_repack(
b_q_weight: torch.Tensor,
perm: torch.Tensor,
python/sglang/jit_kernel/hicache.py
+3
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
@@ -66,6 +67,7 @@ def _default_unroll(element_size: int) -> int:
return 1
@maybe_wrap_jit_kernel_debug
def transfer_hicache_one_layer(
k_cache_dst: torch.Tensor,
v_cache_dst: torch.Tensor,
@@ -101,6 +103,7 @@ def transfer_hicache_one_layer(
)
@maybe_wrap_jit_kernel_debug
def transfer_hicache_all_layer(
k_ptr_dst: torch.Tensor,
v_ptr_dst: torch.Tensor,
python/sglang/jit_kernel/moe_wna16_marlin.py
+2
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
@@ -36,6 +37,7 @@ def _or_empty(
return t if t is not None else torch.empty(0, device=device, dtype=dtype)
@maybe_wrap_jit_kernel_debug
def moe_wna16_marlin_gemm(
a: torch.Tensor,
c_or_none: Optional[torch.Tensor],
python/sglang/jit_kernel/ngram_embedding.py
+3
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
from typing import TYPE_CHECKING
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit
if TYPE_CHECKING:
@@ -21,6 +22,7 @@ def _jit_ngram_embedding_module() -> Module:
)
@maybe_wrap_jit_kernel_debug
def compute_n_gram_ids(
ne_n: int,
ne_k: int,
@@ -66,6 +68,7 @@ def compute_n_gram_ids(
)
@maybe_wrap_jit_kernel_debug
def update_token_table(
tokens: torch.Tensor,
ne_token_table: torch.Tensor,
python/sglang/jit_kernel/norm.py
+6
View File
@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
logger = logging.getLogger(__name__)
from sglang.jit_kernel.utils import (
@@ -78,6 +80,7 @@ def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool:
return False
@maybe_wrap_jit_kernel_debug
def fused_inplace_qknorm(
q: torch.Tensor,
k: torch.Tensor,
@@ -92,6 +95,7 @@ def fused_inplace_qknorm(
module.qknorm(q, k, q_weight, k_weight, eps)
@maybe_wrap_jit_kernel_debug
def rmsnorm(
input: torch.Tensor,
weight: torch.Tensor,
@@ -104,6 +108,7 @@ def rmsnorm(
module.rmsnorm(input, weight, output, eps)
@maybe_wrap_jit_kernel_debug
def fused_add_rmsnorm(
input: torch.Tensor,
residual: torch.Tensor,
@@ -114,6 +119,7 @@ def fused_add_rmsnorm(
module.fused_add_rmsnorm(input, residual, weight, eps)
@maybe_wrap_jit_kernel_debug
def fused_inplace_qknorm_across_heads(
q: torch.Tensor,
k: torch.Tensor,
python/sglang/jit_kernel/nvfp4.py
+7
View File
@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Optional, Tuple
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.srt.utils.custom_op import register_custom_op
@@ -195,6 +196,7 @@ def _jit_nvfp4_blockwise_moe_module() -> Module:
)
@maybe_wrap_jit_kernel_debug
def cutlass_scaled_fp4_mm(
a: torch.Tensor,
b: torch.Tensor,
@@ -211,6 +213,7 @@ def cutlass_scaled_fp4_mm(
return out
@maybe_wrap_jit_kernel_debug
def cutlass_fp4_group_mm(
a_fp4: torch.Tensor,
b_fp4: torch.Tensor,
@@ -290,6 +293,7 @@ def _scaled_fp4_quant_custom_op(
module.scaled_fp4_quant(output, input, output_scale, input_global_scale)
@maybe_wrap_jit_kernel_debug
def scaled_fp4_quant(
input: torch.Tensor, input_global_scale: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -359,6 +363,7 @@ def _scaled_fp4_experts_quant_custom_op(
)
@maybe_wrap_jit_kernel_debug
def scaled_fp4_experts_quant(
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
@@ -443,6 +448,7 @@ def _scaled_fp4_grouped_quant_custom_op(
)
@maybe_wrap_jit_kernel_debug
def scaled_fp4_grouped_quant(
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
@@ -503,6 +509,7 @@ def _silu_and_mul_scaled_fp4_grouped_quant_custom_op(
)
@maybe_wrap_jit_kernel_debug
def silu_and_mul_scaled_fp4_grouped_quant(
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
python/sglang/jit_kernel/per_tensor_quant_fp8.py
+2
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
@@ -22,6 +23,7 @@ def _jit_per_tensor_quant_fp8_module(is_static: bool, dtype: torch.dtype) -> Mod
)
@maybe_wrap_jit_kernel_debug
@register_custom_op(
op_name="per_tensor_quant_fp8",
mutates_args=["output_q", "output_s"],
python/sglang/jit_kernel/per_token_group_quant_8bit.py
+2
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
@@ -72,6 +73,7 @@ def _per_token_group_quant_8bit_custom_op(
return None
@maybe_wrap_jit_kernel_debug
def per_token_group_quant_8bit(
input: torch.Tensor,
output_q: torch.Tensor,
python/sglang/jit_kernel/rope.py
+2
View File
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
@@ -176,6 +177,7 @@ def apply_rope_inplace_with_kvcache(
# NOTE: this name is intentionally set as the old kernel in `sgl_kernel`
@maybe_wrap_jit_kernel_debug
def apply_rope_with_cos_sin_cache_inplace(
q: torch.Tensor,
k: torch.Tensor,
python/sglang/jit_kernel/timestep_embedding.py
+2
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
@@ -21,6 +22,7 @@ def _jit_timestep_embedding_module(dtype: torch.dtype) -> Module:
)
@maybe_wrap_jit_kernel_debug
def timestep_embedding(
t: torch.Tensor,
dim: int,
python/sglang/kernel_api_logging.py
+470
View File
@@ -0,0 +1,470 @@
"""Kernel API crash debugging helpers for SGLang.
This module was developed with reference to FlashInfer's kernel API logging utility:
https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/api_logging.py
"""
from __future__ import annotations
import fnmatch
import functools
import inspect
import json
import logging
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Callable
import torch
def _substitute_process_id(path: str) -> str:
if "%i" in path:
return path.replace("%i", str(os.getpid()))
return path
_KERNEL_API_LOG_LEVEL = int(os.environ.get("SGLANG_KERNEL_API_LOGLEVEL", "0"))
_KERNEL_API_LOG_DEST = _substitute_process_id(
os.environ.get("SGLANG_KERNEL_API_LOGDEST", "stdout")
)
_DUMP_DIR = Path(
_substitute_process_id(
os.environ.get("SGLANG_KERNEL_API_DUMP_DIR", "sglang_kernel_api_dumps")
)
)
_DUMP_INCLUDE_PATTERNS = [
p.strip()
for p in os.environ.get("SGLANG_KERNEL_API_DUMP_INCLUDE", "").split(",")
if p.strip()
]
_DUMP_EXCLUDE_PATTERNS = [
p.strip()
for p in os.environ.get("SGLANG_KERNEL_API_DUMP_EXCLUDE", "").split(",")
if p.strip()
]
_logger = logging.getLogger("sglang.kernel_api")
_dump_call_counter: dict[str, int] = {}
def _setup_logger() -> None:
for handler in list(_logger.handlers):
_logger.removeHandler(handler)
try:
handler.close()
except Exception:
pass
if _KERNEL_API_LOG_LEVEL == 0:
_logger.addHandler(logging.NullHandler())
_logger.setLevel(logging.CRITICAL + 1)
return
_logger.setLevel(logging.DEBUG)
if _KERNEL_API_LOG_DEST == "stdout":
handler = logging.StreamHandler(sys.stdout)
elif _KERNEL_API_LOG_DEST == "stderr":
handler = logging.StreamHandler(sys.stderr)
else:
handler = logging.FileHandler(_KERNEL_API_LOG_DEST, mode="a")
handler.setFormatter(logging.Formatter("%(message)s"))
_logger.addHandler(handler)
_logger.propagate = False
_setup_logger()
def _is_compiling() -> bool:
try:
if hasattr(torch, "compiler") and hasattr(torch.compiler, "is_compiling"):
return bool(torch.compiler.is_compiling())
if hasattr(torch, "_dynamo") and hasattr(torch._dynamo, "is_compiling"):
return bool(torch._dynamo.is_compiling())
except Exception:
return False
return False
def _timestamp() -> str:
return datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")
def _is_cuda_graph_capture_active() -> bool:
try:
return torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()
except Exception:
return False
def _append_line(lines: list[str], indent: int, text: str) -> None:
lines.append(" " * indent + text)
def _should_dump_function(func_name: str) -> bool:
if _DUMP_INCLUDE_PATTERNS and not any(
fnmatch.fnmatch(func_name, pattern) for pattern in _DUMP_INCLUDE_PATTERNS
):
return False
if _DUMP_EXCLUDE_PATTERNS and any(
fnmatch.fnmatch(func_name, pattern) for pattern in _DUMP_EXCLUDE_PATTERNS
):
return False
return True
def _serialize_tensor(tensor: torch.Tensor) -> list[str]:
lines = ["Tensor("]
_append_line(lines, 2, f"shape={tuple(tensor.shape)}")
_append_line(lines, 2, f"dtype={tensor.dtype}")
_append_line(lines, 2, f"device={tensor.device}")
_append_line(lines, 2, f"requires_grad={tensor.requires_grad}")
_append_line(lines, 2, f"is_contiguous={tensor.is_contiguous()}")
if _KERNEL_API_LOG_LEVEL >= 5:
if tensor.numel() == 0:
_append_line(lines, 2, "statistics=[empty tensor]")
elif tensor.device.type == "cuda" and _is_cuda_graph_capture_active():
_append_line(
lines, 2, "statistics=[skipped: CUDA graph capture in progress]"
)
else:
try:
detached = tensor.detach()
if detached.is_complex():
stats_source = detached.abs().float()
nan_count = int(torch.isnan(detached).sum().item())
inf_count = int(torch.isinf(detached).sum().item())
else:
stats_source = detached.float()
if detached.is_floating_point():
nan_count = int(torch.isnan(detached).sum().item())
inf_count = int(torch.isinf(detached).sum().item())
else:
nan_count = 0
inf_count = 0
_append_line(lines, 2, f"min={stats_source.min().item():.6f}")
_append_line(lines, 2, f"max={stats_source.max().item():.6f}")
_append_line(lines, 2, f"mean={stats_source.mean().item():.6f}")
_append_line(lines, 2, f"nan_count={nan_count}")
_append_line(lines, 2, f"inf_count={inf_count}")
except Exception as exc:
_append_line(
lines, 2, f"statistics=[unavailable: {type(exc).__name__}]"
)
lines.append(")")
return lines
def _serialize_value(value: Any, depth: int = 0) -> list[str]:
if depth >= 2:
return [f"{type(value).__name__}(...)"]
if isinstance(value, torch.Tensor):
return _serialize_tensor(value)
if isinstance(value, (str, int, float, bool, type(None))):
return [repr(value)]
if isinstance(value, (list, tuple)):
opener = "[" if isinstance(value, list) else "("
closer = "]" if isinstance(value, list) else ")"
lines = [opener]
for idx, item in enumerate(value[:4]):
item_lines = _serialize_value(item, depth + 1)
lines.append(f" [{idx}] {item_lines[0]}")
for extra in item_lines[1:]:
lines.append(f" {extra}")
if len(value) > 4:
lines.append(f" ... ({len(value) - 4} more items)")
lines.append(closer)
return lines
if isinstance(value, dict):
lines = ["{"]
items = list(value.items())
for key, item in items[:8]:
item_lines = _serialize_value(item, depth + 1)
lines.append(f" {key!r}: {item_lines[0]}")
for extra in item_lines[1:]:
lines.append(f" {extra}")
if len(items) > 8:
lines.append(f" ... ({len(items) - 8} more items)")
lines.append("}")
return lines
summary = [f"{type(value).__name__}("]
for attr in ("shape", "dtype", "device"):
if hasattr(value, attr):
try:
_append_line(summary, 2, f"{attr}={getattr(value, attr)}")
except Exception:
pass
if len(summary) == 1:
_append_line(summary, 2, f"repr={repr(value)[:200]}")
summary.append(")")
return summary
def _serialize_json_value(value: Any) -> Any:
if isinstance(value, torch.dtype):
return {"type": "torch.dtype", "value": str(value)}
if isinstance(value, (str, int, float, bool, type(None))):
return value
if isinstance(value, (list, tuple)):
return [_serialize_json_value(item) for item in value[:16]]
if isinstance(value, dict):
return {
str(key): _serialize_json_value(item)
for key, item in list(value.items())[:32]
}
return {"type": type(value).__name__, "repr": repr(value)[:200]}
def _collect_dump_entries(
prefix: str,
value: Any,
tensor_entries: dict[str, torch.Tensor],
metadata_entries: dict[str, Any],
) -> None:
if isinstance(value, torch.Tensor):
tensor_entries[prefix] = value.detach().cpu()
return
if isinstance(value, (list, tuple)):
for idx, item in enumerate(value):
_collect_dump_entries(
f"{prefix}_{idx}", item, tensor_entries, metadata_entries
)
metadata_entries[f"{prefix}__container"] = {
"type": type(value).__name__,
"length": len(value),
}
return
if isinstance(value, dict):
for key, item in value.items():
_collect_dump_entries(
f"{prefix}_{str(key)}", item, tensor_entries, metadata_entries
)
metadata_entries[f"{prefix}__container"] = {
"type": "dict",
"keys": [str(k) for k in value.keys()],
}
return
metadata_entries[prefix] = _serialize_json_value(value)
def _dump_metadata_path(dump_dir: Path) -> Path:
return dump_dir / "metadata.json"
def _write_dump_metadata(dump_dir: Path, metadata: dict[str, Any]) -> None:
_dump_metadata_path(dump_dir).write_text(json.dumps(metadata, indent=2))
def _read_dump_metadata(dump_dir: Path) -> dict[str, Any]:
return json.loads(_dump_metadata_path(dump_dir).read_text())
def _dump_function_inputs(
func_name: str, args: tuple[Any, ...], kwargs: dict[str, Any]
) -> Path | None:
if not _should_dump_function(func_name):
return None
_DUMP_DIR.mkdir(parents=True, exist_ok=True)
call_index = _dump_call_counter.get(func_name, 0) + 1
_dump_call_counter[func_name] = call_index
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
safe_func_name = func_name.replace("/", "_").replace("<", "_").replace(">", "_")
dump_dir = (
_DUMP_DIR
/ f"{timestamp}_pid{os.getpid()}_{safe_func_name}_call{call_index:04d}"
)
dump_dir.mkdir(parents=True, exist_ok=True)
tensor_entries: dict[str, torch.Tensor] = {}
metadata_entries: dict[str, Any] = {}
for idx, arg in enumerate(args):
_collect_dump_entries(f"arg_{idx}", arg, tensor_entries, metadata_entries)
for key, value in kwargs.items():
_collect_dump_entries(f"kwarg_{key}", value, tensor_entries, metadata_entries)
if tensor_entries:
torch.save(tensor_entries, dump_dir / "inputs.pt")
metadata = {
"function_name": func_name,
"timestamp": timestamp,
"process_id": os.getpid(),
"execution_status": "inputs_saved",
"input_metadata": metadata_entries,
"input_tensor_keys": list(tensor_entries.keys()),
"output_metadata": {},
"output_tensor_keys": [],
}
_write_dump_metadata(dump_dir, metadata)
_logger.debug("Dumped inputs to: %s", dump_dir)
return dump_dir
def _dump_function_outputs(dump_dir: Path, result: Any) -> None:
tensor_entries: dict[str, torch.Tensor] = {}
metadata_entries: dict[str, Any] = {}
_collect_dump_entries("result", result, tensor_entries, metadata_entries)
if tensor_entries:
torch.save(tensor_entries, dump_dir / "outputs.pt")
metadata = _read_dump_metadata(dump_dir)
metadata["execution_status"] = "completed"
metadata["output_metadata"] = metadata_entries
metadata["output_tensor_keys"] = list(tensor_entries.keys())
_write_dump_metadata(dump_dir, metadata)
_logger.debug("Dumped outputs to: %s", dump_dir)
def _mark_dump_exception(dump_dir: Path, exc: Exception) -> None:
metadata = _read_dump_metadata(dump_dir)
metadata["execution_status"] = "exception"
metadata["exception"] = {
"type": type(exc).__name__,
"message": str(exc),
}
_write_dump_metadata(dump_dir, metadata)
def _log_section(title: str, data: dict[str, Any]) -> None:
_logger.debug(title)
for key, value in data.items():
lines = _serialize_value(value)
_logger.debug(" %s=%s", key, lines[0])
for line in lines[1:]:
_logger.debug(" %s", line)
def _infer_func_name(func: Callable) -> str:
qualname = getattr(func, "__qualname__", getattr(func, "__name__", "unknown"))
qualname = qualname.replace(".<locals>.", ".").replace("<locals>.", "")
module = getattr(func, "__module__", "")
for prefix in ("sglang.", "sgl_kernel."):
if module.startswith(prefix):
module = module[len(prefix) :]
break
if module and module not in {"__main__", "builtins"}:
return f"{module}.{qualname}"
source_path = inspect.getsourcefile(func)
if source_path is not None:
return f"{Path(source_path).stem}.{qualname}"
return qualname
def debug_kernel_api(
func: Callable | None = None,
*,
op_name: str | None = None,
) -> Callable:
if _KERNEL_API_LOG_LEVEL == 0:
if func is None:
return lambda f: f
return func
def decorator(f: Callable) -> Callable:
@functools.wraps(f)
def wrapper(*args: Any, **kwargs: Any) -> Any:
if _is_compiling():
return f(*args, **kwargs)
func_name = op_name or _infer_func_name(f)
dump_dir: Path | None = None
positional_args = args
try:
parameters = tuple(inspect.signature(f).parameters.values())
except (TypeError, ValueError):
parameters = ()
if args and parameters and parameters[0].name in {"self", "cls"}:
positional_args = args[1:]
_logger.debug("=" * 80)
_logger.debug("%s SGLang Kernel API Call: %s", _timestamp(), func_name)
if _KERNEL_API_LOG_LEVEL >= 3:
if positional_args:
_log_section(
"Positional input arguments:",
{f"arg[{idx}]": arg for idx, arg in enumerate(positional_args)},
)
if kwargs:
_log_section("Keyword input arguments:", kwargs)
if _KERNEL_API_LOG_LEVEL >= 10:
if _is_cuda_graph_capture_active():
_logger.debug("Tensor dump skipped: CUDA graph capture in progress")
else:
dump_dir = _dump_function_inputs(func_name, positional_args, kwargs)
try:
result = f(*args, **kwargs)
except Exception as exc:
if dump_dir is not None:
_mark_dump_exception(dump_dir, exc)
_logger.debug(
"%s SGLang Kernel API Exception: %s (%s: %s)",
_timestamp(),
func_name,
type(exc).__name__,
exc,
)
raise
if dump_dir is not None:
_dump_function_outputs(dump_dir, result)
if _KERNEL_API_LOG_LEVEL >= 3:
_log_section("Output:", {"return": result})
return result
return wrapper
if func is None:
return decorator
return decorator(func)
def debug_torch_op(op_name: str, *, namespace: str = "sglang") -> Callable:
def call(*args: Any, **kwargs: Any) -> Any:
return getattr(getattr(torch.ops, namespace), op_name)(*args, **kwargs)
return debug_kernel_api(call, op_name=f"{namespace}.custom_op.{op_name}")
def wrap_method_with_debug_kernel_once(
obj: Any,
method_name: str,
*,
op_name: str,
marker_attr: str | None = None,
) -> Any:
if marker_attr is None:
marker_attr = f"_debug_kernel_{method_name}_wrapped"
if getattr(obj, marker_attr, False):
return obj
setattr(
obj,
method_name,
debug_kernel_api(getattr(obj, method_name), op_name=op_name),
)
setattr(obj, marker_attr, True)
return obj
python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py
+9
View File
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
import torch
from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@@ -168,3 +169,11 @@ class AttentionImpl(ABC, Generic[T]):
attn_metadata: T,
) -> torch.Tensor:
raise NotImplementedError
def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
return wrap_method_with_debug_kernel_once(
attn_impl,
"forward",
op_name=f"diffusion.attn_impl.{attn_impl.__class__.__name__}.forward",
)
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+4
View File
@@ -20,6 +20,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
)
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
AttentionImpl,
wrap_attention_impl_forward,
)
from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
from sglang.multimodal_gen.runtime.layers.usp import (
@@ -73,6 +74,7 @@ class UlyssesAttention(nn.Module):
prefix=f"{prefix}.impl",
**extra_impl_args,
)
wrap_attention_impl_forward(self.attn_impl)
self.num_heads = num_heads
self.head_size = head_size
self.num_kv_heads = num_kv_heads
@@ -252,6 +254,7 @@ class LocalAttention(nn.Module):
causal=causal,
**extra_impl_args,
)
wrap_attention_impl_forward(self.attn_impl)
self.num_heads = num_heads
self.head_size = head_size
self.num_kv_heads = num_kv_heads
@@ -338,6 +341,7 @@ class USPAttention(nn.Module):
prefix=f"{prefix}.impl",
**extra_impl_args,
)
wrap_attention_impl_forward(self.attn_impl)
self.num_heads = num_heads
self.head_size = head_size
self.num_kv_heads = num_kv_heads
python/sglang/multimodal_gen/runtime/layers/custom_op.py
+2
View File
@@ -8,6 +8,7 @@ from typing import Any
import torch.nn as nn
from sglang.kernel_api_logging import debug_kernel_api
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -25,6 +26,7 @@ class CustomOp(nn.Module):
super().__init__()
self._forward_method = self.dispatch_forward()
@debug_kernel_api
def forward(self, *args, **kwargs) -> Any:
return self._forward_method(*args, **kwargs)
python/sglang/multimodal_gen/runtime/layers/linear.py
+8
View File
@@ -10,6 +10,7 @@ import torch.distributed as dist
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
from sglang.multimodal_gen.runtime.distributed import (
divide,
get_tp_group,
@@ -195,6 +196,13 @@ class LinearBase(torch.nn.Module):
else:
self.quant_method = quant_config.get_quant_method(self, prefix=prefix)
if self.quant_method is not None:
wrap_method_with_debug_kernel_once(
self.quant_method,
"apply",
op_name=f"diffusion.quant_method.{self.quant_method.__class__.__name__}.apply",
)
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
raise NotImplementedError
python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py
+2
View File
@@ -5,6 +5,7 @@ from typing import Optional, Tuple
import torch
from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
from sglang.kernel_api_logging import debug_kernel_api
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.srt.utils.custom_op import register_custom_op_from_extern
@@ -61,6 +62,7 @@ def _apply_rotary_emb(
return apply_rotary_embedding(x, cos, sin, interleaved)
@debug_kernel_api
def apply_flashinfer_rope_qk_inplace(
q: torch.Tensor,
k: torch.Tensor,
python/sglang/multimodal_gen/runtime/layers/utils.py
+2 -1
View File
@@ -10,6 +10,7 @@ from typing import Any, Callable, List, Optional
import torch
from torch.library import Library
from sglang.kernel_api_logging import debug_torch_op
from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -155,7 +156,7 @@ class CustomOpWrapper:
mutates_args=self.mutates_args,
fake_impl=self.fake_impl,
)
self._impl = getattr(torch.ops.sglang, self.op_name)
self._impl = debug_torch_op(self.op_name)
assert self._impl is not None
return self._impl
python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py
+4
View File
@@ -621,6 +621,9 @@ class SGLangAttentionWrapper(torch.nn.Module):
[nn.Linear(self.inner_dim, query_dim, bias=out_bias), nn.Dropout(dropout)]
)
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
wrap_attention_impl_forward,
)
from sglang.multimodal_gen.runtime.layers.attention.selector import (
get_attn_backend,
)
@@ -636,6 +639,7 @@ class SGLangAttentionWrapper(torch.nn.Module):
num_kv_heads=heads,
causal=False,
)
wrap_attention_impl_forward(self.attn_impl)
self._attn_backend_name = attn_backend.get_enum().name
def forward(
python/sglang/srt/layers/attention/base_attn_backend.py
+2
View File
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.common import is_npu
if TYPE_CHECKING:
@@ -76,6 +77,7 @@ class AttentionBackend(ABC):
"""
raise NotImplementedError()
@debug_kernel_api
def forward(
self,
q: torch.Tensor,
python/sglang/srt/layers/attention/flashinfer_backend.py
+3
View File
@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union
import torch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.environ import envs
@@ -748,6 +749,7 @@ class FlashInferAttnBackend(AttentionBackend):
def get_cuda_graph_seq_len_fill_value(self):
return 1
@debug_kernel_api
def forward_extend(
self,
q: torch.Tensor,
@@ -862,6 +864,7 @@ class FlashInferAttnBackend(AttentionBackend):
return o.view(-1, layer.tp_q_head_num * layer.head_dim)
@debug_kernel_api
def forward_decode(
self,
q: torch.Tensor,
python/sglang/srt/layers/linear.py
+8
View File
@@ -10,6 +10,7 @@ import torch
from torch import nn
from torch.nn.parameter import Parameter, UninitializedParameter
from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
from sglang.srt.distributed import (
divide,
get_tensor_model_parallel_rank,
@@ -176,6 +177,13 @@ class LinearBase(torch.nn.Module):
else:
self.quant_method = quant_config.get_quant_method(self, prefix=prefix)
if self.quant_method is not None:
wrap_method_with_debug_kernel_once(
self.quant_method,
"apply",
op_name=f"sglang.quant_method.{self.quant_method.__class__.__name__}.apply",
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
raise NotImplementedError
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+14 -3
View File
@@ -9,6 +9,7 @@ from torch.nn.parameter import Parameter
# Import to register custom ops for torch.compile compatibility
import sglang.srt.layers.moe.flashinfer_trtllm_moe # noqa: F401
from sglang.kernel_api_logging import debug_torch_op
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
@@ -44,6 +45,16 @@ elif is_cuda_alike():
else:
fp4_quantize = None
_trtllm_fp8_block_scale_routed_moe_wrapper = debug_torch_op(
"trtllm_fp8_block_scale_routed_moe_wrapper"
)
_trtllm_fp8_block_scale_moe_wrapper = debug_torch_op(
"trtllm_fp8_block_scale_moe_wrapper"
)
_trtllm_fp8_per_tensor_scale_moe = debug_torch_op(
"trtllm_fp8_per_tensor_scale_moe_wrapper"
)
def align_fp8_moe_weights_for_flashinfer_trtllm(
layer: Module, swap_w13_halves: bool = False
@@ -375,7 +386,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
topk_weights=topk_output.topk_weights,
)
output = torch.ops.sglang.trtllm_fp8_block_scale_routed_moe_wrapper(
output = _trtllm_fp8_block_scale_routed_moe_wrapper(
topk_ids=packed_topk_ids,
routing_bias=None,
hidden_states=a_q,
@@ -408,7 +419,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
else:
assert TopKOutputChecker.format_is_bypassed(topk_output)
output = torch.ops.sglang.trtllm_fp8_block_scale_moe_wrapper(
output = _trtllm_fp8_block_scale_moe_wrapper(
routing_logits=(
router_logits.to(torch.float32)
if routing_method_type == RoutingMethodType.DeepSeekV3
@@ -465,7 +476,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
# Move kernel call outside context manager to avoid graph breaks
# during torch.compile for piecewise cuda graph.
# Use custom op wrapper for torch.compile compatibility.
output = torch.ops.sglang.trtllm_fp8_per_tensor_scale_moe_wrapper(
output = _trtllm_fp8_per_tensor_scale_moe(
routing_logits=router_logits.to(torch.bfloat16),
routing_bias=routing_bias_cast,
hidden_states=a_q,
python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py
+3
View File
@@ -5,6 +5,7 @@ from typing import NamedTuple, Optional
import torch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
from sglang.srt.layers.moe.token_dispatcher import (
@@ -167,6 +168,7 @@ class FlashinferDispatcher(BaseDispatcher):
(1, self.router_topk), dtype=torch.float32, device="cuda"
)
@debug_kernel_api
def dispatch(
self, hidden_states: torch.Tensor, topk_output: TopKOutput
) -> FlashinferDispatchOutput:
@@ -243,6 +245,7 @@ class FlashinferDispatcher(BaseDispatcher):
moe_output,
)
@debug_kernel_api
def combine(self, combine_input: FlashinferCombineInput) -> torch.Tensor:
hidden_states = combine_input.hidden_states
output_hidden_size = hidden_states.shape[-1]
python/sglang/srt/layers/quantization/bitsandbytes.py
+2 -1
View File
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional
import torch
from packaging import version
from sglang.kernel_api_logging import debug_torch_op
from sglang.srt.layers.linear import LinearBase
from sglang.srt.layers.quantization.base_config import (
FusedMoEMethodBase,
@@ -431,7 +432,7 @@ try:
mutates_args=["out"],
fake_impl=_apply_bnb_4bit_fake,
)
apply_bnb_4bit = torch.ops.sglang.apply_bnb_4bit
apply_bnb_4bit = debug_torch_op("apply_bnb_4bit")
except AttributeError as error:
raise error
python/sglang/srt/layers/quantization/fp8.py
+4 -1
View File
@@ -10,6 +10,7 @@ import torch.nn.functional as F
from torch.nn import Module
from torch.nn.parameter import Parameter
from sglang.kernel_api_logging import debug_torch_op
from sglang.srt.distributed import get_tensor_model_parallel_world_size, get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
@@ -110,6 +111,8 @@ ACTIVATION_SCHEMES = ["static", "dynamic"]
logger = logging.getLogger(__name__)
_apply_fp8_marlin_linear = debug_torch_op("apply_fp8_marlin_linear")
class Fp8Config(QuantizationConfig):
"""Config class for FP8."""
@@ -643,7 +646,7 @@ class Fp8LinearMethod(LinearMethodBase):
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if self.use_marlin:
return torch.ops.sglang.apply_fp8_marlin_linear(
return _apply_fp8_marlin_linear(
input=x,
weight=layer.weight,
weight_scale=layer.weight_scale,
python/sglang/srt/layers/utils/multi_platform.py
+2
View File
@@ -2,6 +2,7 @@ from typing import Callable
from torch import nn
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils import (
cpu_has_amx_support,
is_cpu,
@@ -67,6 +68,7 @@ class MultiPlatformOp(nn.Module):
self.is_torch_compile = False
# Please do not override this method, because `self._forward_method` can change when in torch compile mode
@debug_kernel_api
def forward(self, *args, **kwargs):
return self._forward_method(*args, **kwargs)
python/sglang/srt/mem_cache/memory_pool.py
+5 -1
View File
@@ -37,6 +37,7 @@ import triton
import triton.language as tl
from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.configs.mamba_utils import BaseLinearStateParams
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
@@ -63,7 +64,10 @@ from sglang.srt.utils import (
from sglang.srt.utils.custom_op import register_custom_op
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
store_cache = register_custom_op(store_cache, mutates_args=["k_cache", "v_cache"])
store_cache = register_custom_op(
debug_kernel_api(store_cache, op_name="jit_kernel.kvcache.store_cache"),
mutates_args=["k_cache", "v_cache"],
)
if TYPE_CHECKING:
from sglang.srt.managers.cache_controller import LayerDoneCounter
python/sglang/srt/models/deepseek_common/utils.py
+4 -2
View File
@@ -74,17 +74,19 @@ def awq_dequantize_func():
return awq_dequantize
elif _is_hip:
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.layers.quantization.awq_triton import (
awq_dequantize_triton as awq_dequantize,
)
return awq_dequantize
return debug_kernel_api(awq_dequantize, op_name="DeepseekCommon.awq_dequantize")
elif _is_npu:
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.layers.quantization.awq_triton import (
awq_dequantize_decomposition as awq_dequantize,
)
return awq_dequantize
return debug_kernel_api(awq_dequantize, op_name="DeepseekCommon.awq_dequantize")
else:
return None
python/sglang/srt/models/kimi_vl_moonvit.py
+3
View File
@@ -52,6 +52,8 @@ import torch.nn.functional as F
from transformers.activations import ACT2FN
from transformers.modeling_utils import PreTrainedModel
from sglang.kernel_api_logging import debug_kernel_api
try:
from flash_attn.flash_attn_interface import flash_attn_varlen_func
except ImportError:
@@ -65,6 +67,7 @@ from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.utils import add_prefix
@debug_kernel_api
def multihead_attention(
q: torch.Tensor,
k: torch.Tensor,
python/sglang/srt/models/minimax_m2.py
+3
View File
@@ -25,6 +25,7 @@ import triton.language as tl
from torch import nn
from transformers import PretrainedConfig
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import (
get_moe_expert_parallel_world_size,
@@ -158,6 +159,7 @@ def rmsnorm_apply_kernel_serial(
tl.store(out2_row + offsets2, out2, mask=mask2)
@debug_kernel_api
def rms_sumsq_serial(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
assert x1.is_cuda and x2.is_cuda
B, D1 = x1.shape
@@ -196,6 +198,7 @@ def rms_sumsq_serial(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
return sum_sq
@debug_kernel_api
def rms_apply_serial(
x1: torch.Tensor,
x2: torch.Tensor,
python/sglang/srt/utils/custom_op.py
+4 -2
View File
@@ -6,6 +6,8 @@ from typing import Any, Callable, List, Optional, TypeVar, Union, overload
import torch
import torch.library
from sglang.kernel_api_logging import debug_torch_op
F = TypeVar("F", bound=Callable)
@@ -159,7 +161,7 @@ class CustomOpWrapper:
mutates_args=self.mutates_args,
fake_impl=self.fake_impl,
)
self._impl = getattr(torch.ops.sglang, self.op_name)
self._impl = debug_torch_op(self.op_name)
assert self._impl is not None
return self._impl
@@ -332,4 +334,4 @@ def register_custom_op_from_extern(
fake_impl=fake_impl,
)
return getattr(torch.ops.sglang, name)
return debug_torch_op(name)
sgl-kernel/python/sgl_kernel/__init__.py
+83
View File
@@ -1,4 +1,5 @@
import torch
from sgl_kernel.debug_utils import maybe_wrap_debug_kernel
from sgl_kernel.load_utils import _load_architecture_specific_ops, _preload_cuda_library
# Initialize the ops library based on current GPU
@@ -114,6 +115,88 @@ if torch.version.hip is not None:
from sgl_kernel.elementwise import gelu_quick
_DEBUG_EXPORT_NAMES = [
"apply_shuffle_mul_sum",
"apply_token_bitmask_inplace_cuda",
"awq_dequantize",
"bmm_fp8",
"build_tree_kernel_efficient",
"causal_conv1d_fwd",
"causal_conv1d_update",
"concat_mla_absorb_q",
"concat_mla_k",
"copy_to_gpu_no_ce",
"cutlass_mla_decode",
"cutlass_mla_get_workspace_size",
"downcast_fp8",
"dsv3_fused_a_gemm",
"dsv3_router_gemm",
"es_fp8_blockwise_scaled_grouped_mm",
"es_sm100_mxfp8_blockscaled_grouped_mm",
"es_sm100_mxfp8_blockscaled_grouped_quant",
"fast_topk",
"fast_topk_transform_fused",
"fast_topk_transform_ragged_fused",
"fast_topk_v2",
"fp8_blockwise_scaled_grouped_mm",
"fp8_blockwise_scaled_mm",
"fp8_scaled_mm",
"fused_add_rmsnorm",
"fused_qk_norm_rope",
"gelu_and_mul",
"gelu_tanh_and_mul",
"gemma_fused_add_rmsnorm",
"gemma_rmsnorm",
"gptq_gemm",
"gptq_shuffle",
"int8_scaled_mm",
"kimi_k2_moe_fused_gate",
"merge_state",
"merge_state_v2",
"moe_align_block_size",
"moe_fused_gate",
"moe_sum",
"moe_sum_reduce",
"prepare_moe_input",
"qserve_w4a8_per_chn_gemm",
"qserve_w4a8_per_group_gemm",
"reconstruct_indices_from_tree_mask",
"rmsnorm",
"rotary_embedding",
"segment_packbits",
"sgl_per_token_group_quant_8bit",
"sgl_per_token_group_quant_fp8",
"sgl_per_token_group_quant_int8",
"sgl_per_token_quant_fp8",
"shuffle_rows",
"silu_and_mul",
"top_k_mask_logits",
"top_k_renorm_prob",
"top_p_renorm_prob",
"topk_sigmoid",
"topk_softmax",
"transfer_kv_all_layer",
"transfer_kv_all_layer_mla",
"transfer_kv_per_layer",
"transfer_kv_per_layer_mla",
"tree_speculative_sampling_target_only",
"verify_tree_greedy",
"weak_ref_tensor",
]
if torch.version.hip is not None:
_DEBUG_EXPORT_NAMES.append("gelu_quick")
for _name in _DEBUG_EXPORT_NAMES:
if _name in globals():
globals()[_name] = maybe_wrap_debug_kernel(
globals()[_name], f"sgl_kernel.{_name}"
)
del _name
del _DEBUG_EXPORT_NAMES
def create_greenctx_stream_by_value(*args, **kwargs):
from sgl_kernel.spatial import create_greenctx_stream_by_value as _impl
sgl-kernel/python/sgl_kernel/debug_utils.py
+45
View File
@@ -0,0 +1,45 @@
import os
from typing import Any, Callable, TypeVar, cast, overload
F = TypeVar("F", bound=Callable[..., Any])
def _wrap_debug_kernel(func: F, op_name: str | None = None) -> F:
try:
if int(os.environ.get("SGLANG_KERNEL_API_LOGLEVEL", "0")) == 0:
return func
except Exception:
return func
try:
from sglang.kernel_api_logging import debug_kernel_api
except Exception:
return func
if getattr(func, "_debug_kernel_wrapped", False):
return func
wrapped = debug_kernel_api(func, op_name=op_name)
setattr(wrapped, "_debug_kernel_wrapped", True)
return cast(F, wrapped)
@overload
def maybe_wrap_debug_kernel(func: F) -> F: ...
@overload
def maybe_wrap_debug_kernel(func: F, op_name: str) -> F: ...
@overload
def maybe_wrap_debug_kernel(*, op_name: str | None = None) -> Callable[[F], F]: ...
def maybe_wrap_debug_kernel(
func: F | None = None, op_name: str | None = None
) -> F | Callable[[F], F]:
if func is None:
return lambda wrapped_func: _wrap_debug_kernel(wrapped_func, op_name)
return _wrap_debug_kernel(func, op_name)
sgl-kernel/python/sgl_kernel/flash_attn.py
+3
View File
@@ -2,6 +2,7 @@ from functools import lru_cache
from typing import Optional, Union
import torch
from sgl_kernel.debug_utils import maybe_wrap_debug_kernel
try:
from sgl_kernel import flash_ops
@@ -31,6 +32,7 @@ def maybe_contiguous(x):
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
@maybe_wrap_debug_kernel
def flash_attn_with_kvcache(
q,
k_cache,
@@ -225,6 +227,7 @@ def flash_attn_with_kvcache(
return (out, softmax_lse, *rest) if return_softmax_lse else out
@maybe_wrap_debug_kernel
def flash_attn_varlen_func(
q,
k,