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

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,