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

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
-    rms_norm_fn = rms_norm_fn_native
+    @maybe_wrap_jit_kernel_debug
+    def norm_infer(
+        x: Tensor,
+        weight: Optional[Tensor],
+        bias: Optional[Tensor],
+        eps: float,
+        is_rms_norm: bool = False,
+        out: Optional[Tensor] = None,
+    ):
+        return norm_infer_native(x, weight, bias, eps, is_rms_norm, out)
+
+    @maybe_wrap_jit_kernel_debug
+    def rms_norm_fn(
+        x,
+        weight,
+        bias,
+        residual=None,
+        x1=None,
+        weight1=None,
+        bias1=None,
+        eps=1e-6,
+        dropout_p=0.0,
+        rowscale=None,
+        prenorm=False,
+        residual_in_fp32=False,
+        zero_centered_weight=False,
+        return_dropout_mask=False,
+        out_dtype=None,
+        out=None,
+        residual_out=None,
+    ):
+        return rms_norm_fn_native(
+            x,
+            weight,
+            bias,
+            residual,
+            x1,
+            weight1,
+            bias1,
+            eps,
+            dropout_p,
+            rowscale,
+            prenorm,
+            residual_in_fp32,
+            zero_centered_weight,
+            return_dropout_mask,
+            out_dtype,
+            out,
+            residual_out,
+        )

python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py

@@ -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
-    fuse_scale_shift_gate_select01_kernel = fuse_scale_shift_gate_select01_kernel_native
+    @maybe_wrap_jit_kernel_debug
+    def fuse_scale_shift_kernel(
+        x: torch.Tensor,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+        scale_constant: float = 1.0,
+        block_l: int = 128,
+        block_c: int = 128,
+    ):
+        return fuse_scale_shift_kernel_native(
+            x, scale, shift, scale_constant, block_l, block_c
+        )
+
+    @maybe_wrap_jit_kernel_debug
+    def fuse_scale_shift_gate_select01_kernel(
+        x: torch.Tensor,
+        scale0: torch.Tensor,
+        shift0: torch.Tensor,
+        gate0: torch.Tensor,
+        scale1: torch.Tensor,
+        shift1: torch.Tensor,
+        gate1: torch.Tensor,
+        index: torch.Tensor,
+        block_l: int = 128,
+        block_c: int = 128,
+    ):
+        return fuse_scale_shift_gate_select01_kernel_native(
+            x,
+            scale0,
+            shift0,
+            gate0,
+            scale1,
+            shift1,
+            gate1,
+            index,
+            block_l,
+            block_c,
+        )

python/sglang/jit_kernel/flash_attention_v4.py

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