diff --git a/.claude/skills/debug-cuda-crash/SKILL.md b/.claude/skills/debug-cuda-crash/SKILL.md new file mode 100644 index 000000000..d32126c7a --- /dev/null +++ b/.claude/skills/debug-cuda-crash/SKILL.md @@ -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 | +| | `` | Log to file | +| | `log_%i.txt` | `%i` expands to process ID | +| `SGLANG_KERNEL_API_DUMP_DIR` | `` | 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. diff --git a/docs/diffusion/environment_variables.md b/docs/diffusion/environment_variables.md index 1dec2aafb..c66bf79f6 100644 --- a/docs/diffusion/environment_variables.md +++ b/docs/diffusion/environment_variables.md @@ -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. | diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 1c2f9e0b4..fb4d17797 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -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 diff --git a/python/sglang/jit_kernel/awq_marlin_repack.py b/python/sglang/jit_kernel/awq_marlin_repack.py index 3b06144cb..5c0614629 100644 --- a/python/sglang/jit_kernel/awq_marlin_repack.py +++ b/python/sglang/jit_kernel/awq_marlin_repack.py @@ -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, diff --git a/python/sglang/jit_kernel/debug_utils.py b/python/sglang/jit_kernel/debug_utils.py new file mode 100644 index 000000000..d65ef2feb --- /dev/null +++ b/python/sglang/jit_kernel/debug_utils.py @@ -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) diff --git a/python/sglang/jit_kernel/diffusion/triton/norm.py b/python/sglang/jit_kernel/diffusion/triton/norm.py index 53b1a8255..2c717909b 100644 --- a/python/sglang/jit_kernel/diffusion/triton/norm.py +++ b/python/sglang/jit_kernel/diffusion/triton/norm.py @@ -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, + ) diff --git a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py index 963c59f55..8f3f84939 100644 --- a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py +++ b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py @@ -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) diff --git a/python/sglang/jit_kernel/diffusion/triton/rotary.py b/python/sglang/jit_kernel/diffusion/triton/rotary.py index 011bf7ad0..b8942cd6c 100644 --- a/python/sglang/jit_kernel/diffusion/triton/rotary.py +++ b/python/sglang/jit_kernel/diffusion/triton/rotary.py @@ -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) diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index a34db4d18..8d6761f24 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -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, + ) diff --git a/python/sglang/jit_kernel/flash_attention_v4.py b/python/sglang/jit_kernel/flash_attention_v4.py index dc5e7e4cf..e889cdda3 100644 --- a/python/sglang/jit_kernel/flash_attention_v4.py +++ b/python/sglang/jit_kernel/flash_attention_v4.py @@ -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, diff --git a/python/sglang/jit_kernel/fused_store_index_cache.py b/python/sglang/jit_kernel/fused_store_index_cache.py index dcfbf4585..b1d9af897 100644 --- a/python/sglang/jit_kernel/fused_store_index_cache.py +++ b/python/sglang/jit_kernel/fused_store_index_cache.py @@ -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, diff --git a/python/sglang/jit_kernel/gptq_marlin.py b/python/sglang/jit_kernel/gptq_marlin.py index 8f6047309..8add63f5f 100644 --- a/python/sglang/jit_kernel/gptq_marlin.py +++ b/python/sglang/jit_kernel/gptq_marlin.py @@ -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], diff --git a/python/sglang/jit_kernel/gptq_marlin_repack.py b/python/sglang/jit_kernel/gptq_marlin_repack.py index f04a2ce81..8f7cafc63 100644 --- a/python/sglang/jit_kernel/gptq_marlin_repack.py +++ b/python/sglang/jit_kernel/gptq_marlin_repack.py @@ -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, diff --git a/python/sglang/jit_kernel/hicache.py b/python/sglang/jit_kernel/hicache.py index aa78d5cb4..41861c377 100644 --- a/python/sglang/jit_kernel/hicache.py +++ b/python/sglang/jit_kernel/hicache.py @@ -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, diff --git a/python/sglang/jit_kernel/moe_wna16_marlin.py b/python/sglang/jit_kernel/moe_wna16_marlin.py index 7e17e5ccd..76b2c90dc 100644 --- a/python/sglang/jit_kernel/moe_wna16_marlin.py +++ b/python/sglang/jit_kernel/moe_wna16_marlin.py @@ -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], diff --git a/python/sglang/jit_kernel/ngram_embedding.py b/python/sglang/jit_kernel/ngram_embedding.py index dff20ff64..600880f12 100644 --- a/python/sglang/jit_kernel/ngram_embedding.py +++ b/python/sglang/jit_kernel/ngram_embedding.py @@ -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, diff --git a/python/sglang/jit_kernel/norm.py b/python/sglang/jit_kernel/norm.py index 88a4c0b15..6fad55c44 100644 --- a/python/sglang/jit_kernel/norm.py +++ b/python/sglang/jit_kernel/norm.py @@ -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, diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py index e13c8ac61..7a061a9cc 100644 --- a/python/sglang/jit_kernel/nvfp4.py +++ b/python/sglang/jit_kernel/nvfp4.py @@ -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, diff --git a/python/sglang/jit_kernel/per_tensor_quant_fp8.py b/python/sglang/jit_kernel/per_tensor_quant_fp8.py index 9225aa45d..89cafb7ce 100644 --- a/python/sglang/jit_kernel/per_tensor_quant_fp8.py +++ b/python/sglang/jit_kernel/per_tensor_quant_fp8.py @@ -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"], diff --git a/python/sglang/jit_kernel/per_token_group_quant_8bit.py b/python/sglang/jit_kernel/per_token_group_quant_8bit.py index fdfbb2980..529bb6a99 100644 --- a/python/sglang/jit_kernel/per_token_group_quant_8bit.py +++ b/python/sglang/jit_kernel/per_token_group_quant_8bit.py @@ -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, diff --git a/python/sglang/jit_kernel/rope.py b/python/sglang/jit_kernel/rope.py index d9cbe0a8b..ff9470027 100644 --- a/python/sglang/jit_kernel/rope.py +++ b/python/sglang/jit_kernel/rope.py @@ -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, diff --git a/python/sglang/jit_kernel/timestep_embedding.py b/python/sglang/jit_kernel/timestep_embedding.py index 4f6cc2b42..c0213c5cd 100644 --- a/python/sglang/jit_kernel/timestep_embedding.py +++ b/python/sglang/jit_kernel/timestep_embedding.py @@ -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, diff --git a/python/sglang/kernel_api_logging.py b/python/sglang/kernel_api_logging.py new file mode 100644 index 000000000..661f1f02c --- /dev/null +++ b/python/sglang/kernel_api_logging.py @@ -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("..", ".").replace(".", "") + + 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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 42256d261..b016f0f81 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py @@ -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", + ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 2a695d15a..b60dfd9dc 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/python/sglang/multimodal_gen/runtime/layers/custom_op.py index c4abecb44..30b94beef 100644 --- a/python/sglang/multimodal_gen/runtime/layers/custom_op.py +++ b/python/sglang/multimodal_gen/runtime/layers/custom_op.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index f74309a53..f356fc08d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/linear.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py index 57af17e52..620a4e33b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/layers/utils.py b/python/sglang/multimodal_gen/runtime/layers/utils.py index 2777e9209..a6c7d3b28 100644 --- a/python/sglang/multimodal_gen/runtime/layers/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/utils.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py index 034b7718b..0bb8d2520 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py @@ -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( diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 8d14e32a9..ccfca17e6 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index c51127254..147863803 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -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, diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 8f9febf39..abd756870 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 62005ed46..d32071a39 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -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, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 72d5b2ea3..7b5080bb8 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -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] diff --git a/python/sglang/srt/layers/quantization/bitsandbytes.py b/python/sglang/srt/layers/quantization/bitsandbytes.py index 9f17da1a6..3ee6da386 100644 --- a/python/sglang/srt/layers/quantization/bitsandbytes.py +++ b/python/sglang/srt/layers/quantization/bitsandbytes.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index c7874830d..987162922 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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, diff --git a/python/sglang/srt/layers/utils/multi_platform.py b/python/sglang/srt/layers/utils/multi_platform.py index eff8766e1..ff4d89914 100644 --- a/python/sglang/srt/layers/utils/multi_platform.py +++ b/python/sglang/srt/layers/utils/multi_platform.py @@ -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) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 40fe210df..140d1b514 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index 5be323f29..73f26c2b0 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -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 diff --git a/python/sglang/srt/models/kimi_vl_moonvit.py b/python/sglang/srt/models/kimi_vl_moonvit.py index e41f85695..72f25b6b2 100644 --- a/python/sglang/srt/models/kimi_vl_moonvit.py +++ b/python/sglang/srt/models/kimi_vl_moonvit.py @@ -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, diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 3c7b36496..11929f740 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -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, diff --git a/python/sglang/srt/utils/custom_op.py b/python/sglang/srt/utils/custom_op.py index cd708aa4c..cce713c60 100644 --- a/python/sglang/srt/utils/custom_op.py +++ b/python/sglang/srt/utils/custom_op.py @@ -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) diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index ca56261dd..b5dbca95e 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -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 diff --git a/sgl-kernel/python/sgl_kernel/debug_utils.py b/sgl-kernel/python/sgl_kernel/debug_utils.py new file mode 100644 index 000000000..57e176b15 --- /dev/null +++ b/sgl-kernel/python/sgl_kernel/debug_utils.py @@ -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) diff --git a/sgl-kernel/python/sgl_kernel/flash_attn.py b/sgl-kernel/python/sgl_kernel/flash_attn.py index 4281a5e6b..7aff1af85 100644 --- a/sgl-kernel/python/sgl_kernel/flash_attn.py +++ b/sgl-kernel/python/sgl_kernel/flash_attn.py @@ -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,