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

.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 |
+|  | `<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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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
+    @maybe_wrap_jit_kernel_debug
-    rms_norm_fn = rms_norm_fn_native
+    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

@@ -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

@@ -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

@@ -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
+    @maybe_wrap_jit_kernel_debug
-    fuse_scale_shift_gate_select01_kernel = fuse_scale_shift_gate_select01_kernel_native
+    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,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

@@ -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

@@ -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

@@ -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,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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,