[KDA] Support CuTeDSL KDA decode kernel (#21203)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/layers/attention/linear/gdn_backend.py

@@ -64,7 +64,7 @@ class GDNKernelDispatcher:
             self.decode_kernel = triton_kernel
         elif decode_backend.is_cutedsl():
             if not is_cuda():
-                raise ValueError("CuTe DSL backend requires CUDA")
+                raise ValueError("GDN CuTe DSL backend requires CUDA")
             from sglang.srt.layers.attention.linear.kernels.gdn_cutedsl import (
                 CuteDSLGDNKernel,
             )

python/sglang/srt/layers/attention/linear/kda_backend.py

@@ -3,6 +3,9 @@ from typing import Tuple, Union
 import torch
 
 from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
+from sglang.srt.layers.attention.linear.kernels.kda_cutedsl import (
+    CuteDSLKDAKernel,
+)
 from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
 from sglang.srt.layers.attention.linear.utils import (
     LinearAttnKernelBackend,
@@ -14,7 +17,7 @@ from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
     causal_conv1d_update,
 )
 from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
-from sglang.srt.utils import is_cpu, is_npu
+from sglang.srt.utils import is_cpu, is_cuda, is_npu
 from sglang.srt.utils.common import rank0_log
 
 # KDA always uses the triton causal_conv1d_fn (no CUDA override).
@@ -44,6 +47,10 @@ class KDAKernelDispatcher:
 
         if decode_backend.is_triton():
             self.decode_kernel = triton_kernel
+        elif decode_backend.is_cutedsl():
+            if not is_cuda():
+                raise ValueError("KDA CuTe DSL backend requires CUDA")
+            self.decode_kernel = CuteDSLKDAKernel()
         else:
             raise ValueError(
                 f"Unsupported KDA decode backend: {decode_backend}. "

python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py

@@ -0,0 +1,47 @@
+import torch
+
+from sglang.jit_kernel.cutedsl_kda import cutedsl_fused_sigmoid_gating_kda_update
+from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
+    LinearAttnKernelBase,
+)
+
+
+class CuteDSLKDAKernel(LinearAttnKernelBase):
+    """CuTe DSL kernel for KDA decode (CUDA only)."""
+
+    def decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        *,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        return cutedsl_fused_sigmoid_gating_kda_update(
+            A_log=A_log,
+            dt_bias=dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            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,
+        )
+
+    def extend(self, *args, **kwargs):
+        raise NotImplementedError("CuteDSLKDAKernel only supports decode")
+
+    def target_verify(self, *args, **kwargs):
+        raise NotImplementedError("CuteDSLKDAKernel only supports decode")