[jit-kernel] Add CuTe DSL GDN Decode Kernel

python/pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
   "msgspec",
   "ninja",
   "numpy",
-  "nvidia-cutlass-dsl==4.2.1",
+  "nvidia-cutlass-dsl>=4.3.0",
   "nvidia-ml-py",
   "openai-harmony==0.0.4",
   "openai==2.6.1",

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -1,3 +1,5 @@
+import logging
+import os
 from typing import Optional, Union
 
 import torch
@@ -38,6 +40,43 @@
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import is_cuda, is_npu
 
+logger = logging.getLogger(__name__)
+
+USE_CUTEDSL_GDN_DECODE = os.environ.get("SGLANG_USE_CUTEDSL_GDN_DECODE", "0") == "1"
+_cutedsl_gdn_available = None
+_cutedsl_fused_sigmoid_gating_delta_rule_update = None
+_cutedsl_gdn_logged = False
+
+
+def _check_cutedsl_gdn_available():
+    """Check if CuTe DSL GDN kernel is available and import it."""
+    global _cutedsl_gdn_available, _cutedsl_fused_sigmoid_gating_delta_rule_update
+
+    if _cutedsl_gdn_available is not None:
+        return _cutedsl_gdn_available
+
+    try:
+        from sgl_kernel.cutedsl_gdn import (
+            cutedsl_fused_sigmoid_gating_delta_rule_update,
+            is_cutedsl_gdn_available,
+        )
+
+        if is_cutedsl_gdn_available():
+            _cutedsl_fused_sigmoid_gating_delta_rule_update = (
+                cutedsl_fused_sigmoid_gating_delta_rule_update
+            )
+            _cutedsl_gdn_available = True
+            logger.info("CuTe DSL GDN decode kernel is available")
+        else:
+            _cutedsl_gdn_available = False
+            logger.info("CuTe DSL GDN decode kernel: cutlass/cute not installed")
+    except ImportError as e:
+        _cutedsl_gdn_available = False
+        logger.debug(f"CuTe DSL GDN decode kernel not available: {e}")
+
+    return _cutedsl_gdn_available
+
+
 if is_cuda():
     from sglang.srt.layers.attention.mamba.causal_conv1d import (
         causal_conv1d_fn as causal_conv1d_fn_cuda,
@@ -897,21 +936,44 @@ def forward_decode(
         key = key.view(1, seq_len, num_heads, head_k_dim)
         value = value.view(1, seq_len, value.shape[1] // head_v_dim, head_v_dim)
 
-        core_attn_out = fused_sigmoid_gating_delta_rule_update(
-            A_log=A_log,
-            dt_bias=dt_bias,
-            q=query,
-            k=key,
-            v=value,
-            a=a,
-            b=b,
-            initial_state_source=ssm_states,
-            initial_state_indices=cache_indices,
-            cu_seqlens=query_start_loc,
-            use_qk_l2norm_in_kernel=True,
-            softplus_beta=1.0,
-            softplus_threshold=20.0,
-        )
+        global _cutedsl_gdn_logged
+        use_cutedsl = USE_CUTEDSL_GDN_DECODE and _check_cutedsl_gdn_available()
+        if not _cutedsl_gdn_logged and USE_CUTEDSL_GDN_DECODE:
+            logger.info(f"Using CuTe DSL GDN kernel for decode: {use_cutedsl}")
+            _cutedsl_gdn_logged = True
+
+        if use_cutedsl:
+            core_attn_out = _cutedsl_fused_sigmoid_gating_delta_rule_update(
+                A_log=A_log,
+                dt_bias=dt_bias,
+                q=query,
+                k=key,
+                v=value,
+                a=a,
+                b=b,
+                initial_state_source=ssm_states,
+                initial_state_indices=cache_indices,
+                cu_seqlens=query_start_loc,
+                use_qk_l2norm_in_kernel=True,
+                softplus_beta=1.0,
+                softplus_threshold=20.0,
+            )
+        else:
+            core_attn_out = fused_sigmoid_gating_delta_rule_update(
+                A_log=A_log,
+                dt_bias=dt_bias,
+                q=query,
+                k=key,
+                v=value,
+                a=a,
+                b=b,
+                initial_state_source=ssm_states,
+                initial_state_indices=cache_indices,
+                cu_seqlens=query_start_loc,
+                use_qk_l2norm_in_kernel=True,
+                softplus_beta=1.0,
+                softplus_threshold=20.0,
+            )
 
         self._track_mamba_state_decode(
             forward_batch, conv_states, ssm_states, cache_indices

sgl-kernel/python/sgl_kernel/__init__.py

@@ -132,6 +132,16 @@
 )
 from sgl_kernel.version import __version__
 
+# CuTe DSL GDN kernel (optional - requires cutlass/cute)
+try:
+    from sgl_kernel.cutedsl_gdn import (
+        cutedsl_fused_sigmoid_gating_delta_rule_update,
+        is_cutedsl_gdn_available,
+    )
+except ImportError:
+    # cutlass/cute not available
+    pass
+
 if torch.version.hip is not None:
     from sgl_kernel.elementwise import gelu_quick