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

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,