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

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)