[KDA] Support Helion MTP verification and harden prefill, decode

python/sglang/kernels/ops/attention/helion/kda_verify.py

@@ -0,0 +1,145 @@
+"""Helion KDA verification with per-candidate state snapshots."""
+
+from __future__ import annotations
+
+import helion
+import helion.language as hl
+import torch
+
+
+@helion.kernel(
+    static_shapes=False,
+    config=helion.Config(
+        block_sizes=[8],
+        num_warps=1,
+        num_stages=1,
+        indexing="pointer",
+    ),
+    ignore_warnings=[helion.exc.ProcessGroupNameNotFound],
+)
+def _kda_verify(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    state_pool: torch.Tensor,
+    state_indices: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    snapshots: torch.Tensor,
+    snapshot_indices: torch.Tensor,
+    parents: torch.Tensor | None,
+    out: torch.Tensor,
+    lower_bound: float,
+    bounded_gate: hl.constexpr,
+) -> torch.Tensor:
+    N = cu_seqlens.size(0) - 1
+    H = hl.specialize(q.size(1))
+    HV = hl.specialize(v.size(1))
+    K = hl.specialize(q.size(2))
+    V = hl.specialize(v.size(2))
+    block_v = hl.register_block_size(1, V)
+    for tile_n, tile_h, tile_v in hl.tile([N, HV, V], block_size=[1, 1, block_v]):
+        n, hv = tile_n.id, tile_h.id
+        h = hv // (HV // H)
+        keys = hl.arange(K)
+        slot = state_indices[n].long()
+        scratch = snapshot_indices[n].long()
+        start = cu_seqlens[n].long()
+        end = cu_seqlens[n + 1].long()
+        state = hl.load(
+            state_pool, [slot, hv, tile_v.index, keys], extra_mask=slot >= 0
+        ).float()
+        for token in range(start, end):
+            step = token - start
+            if parents is not None:
+                use_parent = (step > 0) & (scratch >= 0)
+                parent = parents[n, step].long()
+                parent_state = hl.load(
+                    snapshots,
+                    [scratch, parent, hv, tile_v.index, keys],
+                    extra_mask=use_parent,
+                ).float()
+                state = torch.where(use_parent, parent_state, state)
+            raw_gate = a[token, hv, keys].float() + dt_bias[hv, keys].float()
+            A = torch.exp(A_log[hv].float())
+            if bounded_gate:
+                log_decay = lower_bound * torch.sigmoid(A * raw_gate)
+            else:
+                softplus = torch.where(
+                    raw_gate <= 20.0, torch.log(1.0 + torch.exp(raw_gate)), raw_gate
+                )
+                log_decay = -A * softplus
+            beta = torch.sigmoid(b[token, hv].float())
+            query = q[token, h, keys].float()
+            key = k[token, h, keys].float()
+            query = query / torch.sqrt((query * query).sum() + 1e-6)
+            key = key / torch.sqrt((key * key).sum() + 1e-6)
+            state = state * torch.exp(log_decay)[None, :]
+            value = v[token, hv, tile_v].float()
+            residual = (value - (state * key[None, :]).sum(-1)) * beta
+            state = state + residual[:, None] * key[None, :]
+            out[token, hv, tile_v] = (state * (query * K**-0.5)[None, :]).sum(-1)
+            if scratch >= 0:
+                snapshots[scratch, step, hv, tile_v.index, keys] = state
+    return out
+
+
+def helion_kda_verify(
+    *,
+    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,
+    intermediate_states_buffer: torch.Tensor,
+    intermediate_state_indices: torch.Tensor,
+    cache_steps: int,
+    retrieve_parent_token: torch.Tensor | None = None,
+    lower_bound: float | None = None,
+) -> torch.Tensor:
+    """Verify packed sequences without mutating the committed state pool.
+
+    Chain verification keeps an FP32 accumulator across candidates; tree
+    verification reads each parent's snapshot, matching the reference contract.
+    Negative state/snapshot indices are CUDA Graph padding slots.
+    """
+    if q.ndim != 4 or q.shape[0] != 1:
+        raise ValueError("Helion KDA verify expects packed [1, T, H, K] inputs")
+    tokens, heads, key_dim = q.shape[1:]
+    value_heads, value_dim = v.shape[2:]
+    if value_heads % heads or k.shape != q.shape:
+        raise ValueError("Inconsistent Helion KDA verify query/key head layout")
+    if intermediate_states_buffer is None or cache_steps <= 0:
+        raise ValueError("Helion KDA verify requires per-candidate state snapshots")
+    if intermediate_states_buffer.shape[1] < cache_steps:
+        raise ValueError(
+            "Helion KDA verify snapshot capacity is smaller than cache_steps"
+        )
+    out = torch.zeros_like(v)
+    _kda_verify(
+        q[0],
+        k[0],
+        v[0],
+        a.reshape(tokens, value_heads, key_dim),
+        b.reshape(tokens, value_heads),
+        A_log.reshape(-1),
+        dt_bias.reshape(value_heads, key_dim),
+        ssm_states,
+        cache_indices,
+        query_start_loc,
+        intermediate_states_buffer,
+        intermediate_state_indices,
+        retrieve_parent_token,
+        out[0],
+        0.0 if lower_bound is None else lower_bound,
+        lower_bound is not None,
+    )
+    return out

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

@@ -58,7 +58,11 @@ def __init__(
         triton_kernel = TritonKDAKernel()
         self.triton_kernel = triton_kernel
         helion_kernel = None
-        if decode_backend.is_helion() or prefill_backend.is_helion():
+        if (
+            decode_backend.is_helion()
+            or prefill_backend.is_helion()
+            or verify_backend.is_helion()
+        ):
             if not is_cuda():
                 raise ValueError("KDA Helion backend requires CUDA")
             from sglang.srt.layers.attention.linear.kernels.kda_helion import (
@@ -106,9 +110,13 @@ def __init__(
         #     the KDA correctness tests assert against.
         #   flashinfer: recurrent_kda (SM100, chain only); reuses the decode kernel
         #     when decode is also flashinfer.
+        #   helion: recurrent chain + tree verify with per-candidate snapshots.
         #   nv_cutedsl: fused Kimi-K3/DSpARK dense verify.
         if verify_backend.is_triton():
             self.verify_kernel = triton_kernel
+        elif verify_backend.is_helion():
+            assert helion_kernel is not None
+            self.verify_kernel = helion_kernel
         elif verify_backend.is_flashinfer():
             if decode_backend.is_flashinfer():
                 self.verify_kernel = self.decode_kernel
@@ -131,7 +139,7 @@ def __init__(
         else:
             raise ValueError(
                 f"Unsupported KDA verify backend: {verify_backend}. "
-                "KDA verify supports 'triton', 'nv_cutedsl', or 'flashinfer'."
+                "KDA verify supports 'triton', 'helion', 'nv_cutedsl', or 'flashinfer'."
             )
 
         if prefill_backend.is_triton():
@@ -203,6 +211,9 @@ def __init__(
         self.supports_packed_decode = getattr(
             self.decode_kernel, "supports_packed_decode", False
         )
+        self.supports_packed_decode_safe_gate = getattr(
+            self.decode_kernel, "supports_packed_decode_safe_gate", False
+        )
 
         rank0_log(
             f"KDA kernel dispatcher: decode={self.decode_kernel.__class__.__name__}, "
@@ -302,15 +313,15 @@ def target_verify(
         lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
-        """MTP / speculative-decode verify, routed to ``self.verify_kernel``
-        (FlashInfer decode -> recurrent_kda; Triton / CuTe DSL decode -> the Triton
-        fused KDA verify)."""
-        if lower_bound is not None and not isinstance(
-            self.verify_kernel, TritonKDAKernel
+        """MTP verification routed to the explicitly selected verify backend."""
+        if (
+            lower_bound is not None
+            and not isinstance(self.verify_kernel, TritonKDAKernel)
+            and not getattr(self.verify_kernel, "supports_safe_gate_verify", False)
         ):
             raise NotImplementedError(
-                "lower_bound (safe gate) target verify is only supported by "
-                f"TritonKDAKernel; got {self.verify_kernel.__class__.__name__}."
+                "lower_bound (safe gate) target verify is not supported by "
+                f"{self.verify_kernel.__class__.__name__}."
             )
         return self.verify_kernel.target_verify(
             A_log=A_log,
@@ -738,9 +749,9 @@ def forward_decode(
         )
 
         # The packed kernel assumes one token per request.
-        if (
-            self.kernel_dispatcher.supports_packed_decode
-            and getattr(layer, "lower_bound", None) is None
+        if self.kernel_dispatcher.supports_packed_decode and (
+            getattr(layer, "lower_bound", None) is None
+            or self.kernel_dispatcher.supports_packed_decode_safe_gate
         ):
             assert qkv.shape[0] == cache_indices.shape[0], (
                 "KDA packed decode requires one token per sequence (T=1): "

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

@@ -11,15 +11,17 @@
 
 
 class HelionKDAKernel(LinearAttnKernelBase):
-    """KDA packed decode and prefill implemented with Helion kernels.
+    """KDA packed decode, prefill, and verification with Helion kernels.
 
-    The generic decode interface delegates to Triton, and the dispatcher routes
-    speculative target verification directly to Triton. The one-token decode
+    The generic decode interface delegates to Triton. The one-token decode
     and ReplaySSM paths use :meth:`packed_decode`, while prefill uses
-    :meth:`extend`.
+    :meth:`extend`. Verification writes per-candidate snapshots without
+    updating the committed state pool.
     """
 
     supports_packed_decode = True
+    supports_packed_decode_safe_gate = True
+    supports_safe_gate_verify = True
     supports_track_state_snapshot: bool = True
 
     def __init__(
@@ -157,6 +159,51 @@ def decode(
             **kwargs,
         )
 
+    def target_verify(
+        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,
+        intermediate_states_buffer: torch.Tensor,
+        intermediate_state_indices: torch.Tensor,
+        cache_steps: int,
+        retrieve_parent_token: torch.Tensor | None = None,
+        lower_bound: float | None = None,
+        cache_ring: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        if cache_ring:
+            raise NotImplementedError(
+                "Helion KDA verify does not support ReplaySSM rings"
+            )
+        from sglang.kernels.ops.attention.helion.kda_verify import helion_kda_verify
+
+        return helion_kda_verify(
+            q=q,
+            k=k,
+            v=v,
+            a=a,
+            b=b,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            intermediate_states_buffer=intermediate_states_buffer,
+            intermediate_state_indices=intermediate_state_indices,
+            cache_steps=cache_steps,
+            retrieve_parent_token=retrieve_parent_token,
+            lower_bound=lower_bound,
+        )
+
     def extend(
         self,
         q: torch.Tensor,
@@ -172,6 +219,7 @@ def extend(
         dt_bias: torch.Tensor | None = None,
         lower_bound: float | None = None,
         return_intermediate_states: bool = False,
+        beta_is_raw: bool = False,
         **kwargs,
     ) -> torch.Tensor:
         assert self._chunk_kda is not None
@@ -188,6 +236,7 @@ def extend(
             A_log=A_log,
             dt_bias=dt_bias,
             lower_bound=lower_bound,
+            beta_is_raw=beta_is_raw,
             output_intermediate_states=return_intermediate_states,
             track_state=kwargs.get("track_state"),
             track_chunk_idx=kwargs.get("track_chunk_idx"),

test/registered/kernels/ops/attention/test_kda_helion.py

@@ -1178,5 +1178,70 @@ def test_packed_varlen_prefill_contract(
     )
 
 
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("lower_bound", [-5.0, None])
+@pytest.mark.parametrize("tree", [False, True])
+def test_target_verify_contract(dtype, lower_bound, tree):
+    from sglang.srt.layers.attention.linear.kda_backend import KDAKernelDispatcher
+    from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
+    from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend
+
+    torch.manual_seed(42)
+    n, t, h, hv, k, v = 4, 6, 2, 4, 128, 128
+    total = 15
+    qkv = torch.randn(1, total, 2 * h * k + hv * v, device="cuda", dtype=torch.bfloat16)
+    q, keys, values = qkv.split([h * k, h * k, hv * v], -1)
+    q = q.view(1, total, h, k)
+    keys = keys.view(1, total, h, k)
+    values = values.view(1, total, hv, v)
+    state = torch.randn(8, hv, v, k, device="cuda", dtype=dtype) * 0.01
+    initial = state.clone()
+    snapshots = torch.full((8, t + 1, hv, v, k), -99.0, device="cuda", dtype=dtype)
+    actual_snapshots = snapshots.clone()
+    args = dict(
+        q=q,
+        k=keys,
+        v=values,
+        a=torch.randn(total, hv * k, device="cuda", dtype=torch.bfloat16),
+        b=torch.randn(total, hv, device="cuda", dtype=torch.bfloat16),
+        A_log=torch.randn(hv, device="cuda"),
+        dt_bias=torch.randn(hv * k, device="cuda"),
+        ssm_states=state,
+        cache_indices=torch.tensor([3, 5, 0, -1], device="cuda", dtype=torch.int32),
+        query_start_loc=torch.tensor(
+            [0, 6, 9, 9, 15], device="cuda", dtype=torch.int32
+        ),
+        intermediate_states_buffer=snapshots,
+        intermediate_state_indices=torch.tensor(
+            [2, 4, 1, -1], device="cuda", dtype=torch.int64
+        ),
+        cache_steps=t,
+        retrieve_parent_token=torch.tensor(
+            [[0, 0, 0, 1, 1, 3]] * n, device="cuda", dtype=torch.int64
+        )
+        if tree
+        else None,
+        lower_bound=lower_bound,
+    )
+    dispatcher = KDAKernelDispatcher(
+        decode_backend=LinearAttnKernelBackend.TRITON,
+        prefill_backend=LinearAttnKernelBackend.TRITON,
+        verify_backend=LinearAttnKernelBackend.HELION,
+    )
+    expected = TritonKDAKernel().target_verify(**args)
+    args["intermediate_states_buffer"] = actual_snapshots
+    actual = dispatcher.target_verify(**args)
+    torch.testing.assert_close(actual, expected, atol=0.003, rtol=0.02)
+    torch.testing.assert_close(actual_snapshots, snapshots, atol=0.003, rtol=0.02)
+    torch.testing.assert_close(state, initial, atol=0, rtol=0)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = dispatcher.target_verify(**args)
+    graph.replay()
+    torch.cuda.synchronize()
+    torch.testing.assert_close(captured, actual, atol=0, rtol=0)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))