[Kernel] Enable Helion backend for Kimi Delta-Attention (#32593)

Co-authored-by: Ethan Che <eche@meta.com>

benchmark/bench_linear_attention/bench_kda_decode.py

@@ -4,6 +4,7 @@ Benchmark & Correctness: KDA Packed Decode vs Baseline Decode.
 Compares:
   - Baseline: split(mixed_qkv) -> view -> fused_sigmoid_gating_delta_rule_update(is_kda=True)
   - Packed:   fused_recurrent_kda_packed_decode (single fused kernel)
+  - Helion:   helion_fused_recurrent_kda_packed_decode
 
 Differences from the GDN packed decode benchmark:
   - KDA gate ``a`` is per-K with shape ``[B, HV * K]`` (instead of ``[B, HV]``).
@@ -28,6 +29,9 @@ from sglang.kernels.ops.attention.fla.fused_recurrent import (
 from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )
+from sglang.kernels.ops.attention.helion.kda_decode import (
+    helion_fused_recurrent_kda_packed_decode,
+)
 
 
 def make_inputs(
@@ -131,6 +135,27 @@ def run_packed(inp):
     return out.transpose(0, 1), ssm_states
 
 
+def run_helion(inp):
+    """Helion path: same packed decode contract and output layout."""
+    B, HV, V = inp["B"], inp["HV"], inp["V"]
+    ssm_states = inp["ssm_states"].clone()
+    out = inp["mixed_qkv"].new_empty(B, 1, HV, V)
+
+    helion_fused_recurrent_kda_packed_decode(
+        mixed_qkv=inp["mixed_qkv"],
+        a=inp["a"],
+        b=inp["b"],
+        A_log=inp["A_log"],
+        dt_bias=inp["dt_bias"],
+        scale=inp["K"] ** -0.5,
+        initial_state=ssm_states,
+        out=out,
+        ssm_state_indices=inp["cache_indices"],
+        use_qk_l2norm_in_kernel=True,
+    )
+    return out.transpose(0, 1), ssm_states
+
+
 def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
     """Run correctness check for a single config. Returns True if PASS."""
     tag = f"B={B:>4} H={H:>2} HV={HV:>2} K={K:>3} V={V:>3} pool={pool_size:>4}"
@@ -139,6 +164,7 @@ def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
 
     o_baseline, state_baseline = run_baseline(inp)
     o_packed, state_packed = run_packed(inp)
+    o_helion, state_helion = run_helion(inp)
 
     atol = 2e-2 if dtype != torch.float32 else 1e-4
     rtol = 1e-2 if dtype != torch.float32 else 1e-4
@@ -157,6 +183,19 @@ def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
         atol, rtol * state_baseline[indices].float().abs().max().item()
     )
 
+    helion_out_diff = (o_helion.float() - o_packed.float()).abs().max().item()
+    helion_state_diff = (
+        (state_helion[indices].float() - state_packed[indices].float())
+        .abs()
+        .max()
+        .item()
+    )
+    helion_ok = helion_out_diff <= max(
+        atol, rtol * o_packed.float().abs().max().item()
+    ) and helion_state_diff <= max(
+        atol, rtol * state_packed[indices].float().abs().max().item()
+    )
+
     passed = output_ok and state_ok
     if passed:
         print(
@@ -166,7 +205,12 @@ def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
         print(
             f"  [FAIL] {tag}  out max_diff={out_diff:.6f}, state max_diff={st_diff:.6f}"
         )
-    return passed
+    print(
+        f"  [{'PASS' if helion_ok else 'FAIL'}] Helion vs packed {tag}  "
+        f"(out max_diff={helion_out_diff:.2e}, "
+        f"state max_diff={helion_state_diff:.2e})"
+    )
+    return passed and helion_ok
 
 
 def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
@@ -213,6 +257,20 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
             use_qk_l2norm_in_kernel=True,
         )
 
+    def fn_helion():
+        helion_fused_recurrent_kda_packed_decode(
+            mixed_qkv=inp["mixed_qkv"],
+            a=inp["a"],
+            b=inp["b"],
+            A_log=inp["A_log"],
+            dt_bias=inp["dt_bias"],
+            scale=K**-0.5,
+            initial_state=inp["ssm_states"],
+            out=out_buf,
+            ssm_state_indices=inp["cache_indices"],
+            use_qk_l2norm_in_kernel=True,
+        )
+
     # Intentionally wall-clock CUDA-event timing, not the shared do_bench /
     # do_bench_cudagraph util: ~2/3 of the packed win is eager CPU dispatch
     # (split + 3x unflatten + extra launch), which graph capture / L2-flush
@@ -222,6 +280,7 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
     for _ in range(warmup):
         fn_baseline()
         fn_packed()
+        fn_helion()
     torch.cuda.synchronize()
 
     def _time(fn):
@@ -236,7 +295,7 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
 
     ms_baseline = _time(fn_baseline)
     ms_packed = _time(fn_packed)
-
+    ms_helion = _time(fn_helion)
     speedup = ms_baseline / ms_packed if ms_packed > 0 else float("inf")
     saved_us = (ms_baseline - ms_packed) * 1000
 
@@ -245,7 +304,9 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
         f"{ms_baseline * 1000:>10.1f} | "
         f"{ms_packed * 1000:>10.1f} | "
         f"{speedup:>7.2f}x | "
-        f"{saved_us:>+9.1f}"
+        f"{saved_us:>+9.1f} | "
+        f"{ms_helion * 1000:>10.1f} | "
+        f"{ms_packed / ms_helion:>10.2f}x"
     )
 
 
@@ -289,8 +350,10 @@ def run_correctness(device, dtype):
     )
     o_baseline, _ = run_baseline(inp)
     o_packed, _ = run_packed(inp)
+    o_helion, _ = run_helion(inp)
     try:
         torch.testing.assert_close(o_packed, o_baseline, atol=2e-2, rtol=1e-2)
+        torch.testing.assert_close(o_helion, o_packed, atol=2e-2, rtol=1e-2)
         print("  [PASS] PAD_SLOT_ID=-1 handling")
     except AssertionError as e:
         print(f"  [FAIL] PAD_SLOT_ID=-1 handling: {e}")
@@ -323,9 +386,11 @@ def run_benchmark(device, dtype, args):
         f"{'base (us)':>10} | "
         f"{'packed (us)':>10} | "
         f"{'speedup':>8} | "
-        f"{'saved (us)':>10}"
+        f"{'saved (us)':>10} | "
+        f"{'Helion (us)':>10} | "
+        f"{'packed/H':>11}"
     )
-    print("  " + "-" * 80)
+    print("  " + "-" * 116)
 
     for B, H, HV in bench_configs:
         # Packed kernel requires HV % H == 0 (GVA / grouped query layout).

benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py

@@ -28,6 +28,9 @@ import argparse
 
 import torch
 
+from sglang.kernels.ops.attention.helion.kda_decode import (
+    helion_fused_recurrent_kda_packed_decode,
+)
 from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
 
 
@@ -59,6 +62,9 @@ def make_decode_inputs(B, H, HV, K, V, pool_size, device, dtype, seed=42):
     ssm = torch.randn(pool_size, HV, V, K, device=device, dtype=dtype) * 0.01
     cache_indices = torch.arange(B, device=device, dtype=torch.int32)
     qsl = torch.arange(B + 1, device=device, dtype=torch.int32)
+    mixed_qkv = torch.cat(
+        (q.reshape(B, H * K), k.reshape(B, H * K), v.reshape(B, HV * V)), dim=-1
+    )
     return dict(
         q=q.contiguous(),
         k=k.contiguous(),
@@ -70,6 +76,7 @@ def make_decode_inputs(B, H, HV, K, V, pool_size, device, dtype, seed=42):
         ssm=ssm.contiguous(),
         cache_indices=cache_indices,
         qsl=qsl,
+        mixed_qkv=mixed_qkv.contiguous(),
         B=B,
         H=H,
         HV=HV,
@@ -139,6 +146,23 @@ def call_decode(kernel, inp, ssm):
     return out.reshape(inp["B"], inp["HV"], inp["V"]).float()
 
 
+def call_helion(inp, ssm):
+    out = inp["mixed_qkv"].new_empty(inp["B"], 1, inp["HV"], inp["V"])
+    helion_fused_recurrent_kda_packed_decode(
+        mixed_qkv=inp["mixed_qkv"],
+        a=inp["a"],
+        b=inp["b"],
+        A_log=inp["A_log"],
+        dt_bias=inp["dt_bias"],
+        scale=inp["K"] ** -0.5,
+        initial_state=ssm,
+        out=out,
+        ssm_state_indices=inp["cache_indices"],
+        use_qk_l2norm_in_kernel=True,
+    )
+    return out.reshape(inp["B"], inp["HV"], inp["V"]).float()
+
+
 def call_verify(kernel, inp, ssm, intermediate_states):
     out = kernel.target_verify(
         A_log=inp["A_log"],
@@ -186,11 +210,14 @@ def run(task, fi, tri, device, dtype, args):
     )
     print("=" * 92)
     hdr = "B" if not is_verify else "B(xT)"
-    print(
+    header = (
         f"  {hdr:>6}  {'H':>3}  {'HV':>3} | {'triton(us)':>11} | "
         f"{'flashinfer(us)':>14} | {'speedup':>8} | {'out_max_diff':>12}"
     )
-    print("  " + "-" * 86)
+    if not is_verify:
+        header += f" | {'helion(us)':>11} | {'FI/H':>8} | {'H max diff':>10}"
+    print(header)
+    print("  " + "-" * (124 if not is_verify else 86))
 
     for B in args.batch_sizes:
         for H in args.num_q_heads:
@@ -235,10 +262,23 @@ def run(task, fi, tri, device, dtype, args):
                 )
                 fi_us = f"{ms_fi * 1000:>14.1f}" if fi is not None else f"{'skip':>14}"
                 sp = f"{speed:>7.2f}x" if fi is not None else f"{'-':>8}"
-                print(
+                line = (
                     f"  {B:>6}  {H:>3}  {HV:>3} | {ms_tri * 1000:>11.1f} | "
                     f"{fi_us} | {sp} | {diff:>12}"
                 )
+                if not is_verify:
+                    helion_state = inp["ssm"].clone()
+                    o_helion = call_helion(inp, helion_state)
+                    helion_diff = (o_helion - o_tri).abs().max().item()
+                    ms_helion = _time(lambda: call_helion(inp, helion_state))
+                    fi_helion = (
+                        f"{ms_fi / ms_helion:>7.2f}x" if fi is not None else f"{'-':>8}"
+                    )
+                    line += (
+                        f" | {ms_helion * 1000:>11.1f} | {fi_helion} | "
+                        f"{helion_diff:>10.2e}"
+                    )
+                print(line)
 
 
 def main():

benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py

@@ -4,6 +4,7 @@ Benchmark & Correctness: Triton KDA vs CuTeDSL KDA (prefill, SM100 Blackwell).
 Compares:
   - Triton:  sglang's chunk_kda (FLA chunkwise gated delta rule, per-channel gate)
   - CuteDSL: kda_blackwell pipeline (fused Triton prologue -> kkt_inv_uw -> h -> o)
+  - Helion:  sglang's Helion chunk_kda
 
 KDA differs from GDN by a PER-CHANNEL decay gate (g is [T, H, K], not scalar).
 The cutedsl pipeline externalizes the per-channel decay into five pre-scaled
@@ -30,6 +31,9 @@ import torch
 import torch.nn.functional as F
 
 from sglang.kernels.ops.attention.fla.kda import chunk_kda, fused_recurrent_kda
+from sglang.kernels.ops.attention.helion.kda_prefill import (
+    chunk_kda as helion_chunk_kda,
+)
 from sglang.kernels.ops.attention.linear.kda_blackwell import prepare_metadata
 from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_h import (
     kda_h_cutedsl,
@@ -132,6 +136,7 @@ def cutedsl_buffers(inp, num_sms, device):
         total=total,
         num_sms=num_sms,
         h0=torch.zeros(1, H, V, K, device=device, dtype=torch.float32),
+        state_indices=torch.zeros(1, device=device, dtype=torch.int32),
         U=torch.empty(pad_t, H, V, device=device, dtype=torch.bfloat16),
         W=torch.empty(pad_t, H, K, device=device, dtype=torch.bfloat16),
         V_new=torch.empty(pad_t, H, V, device=device, dtype=torch.bfloat16),
@@ -172,6 +177,7 @@ def run_cutedsl_pipeline(inp, buf, scale):
         buf["ht"],
         buf["cu"],
         buf["co"],
+        buf["state_indices"],
     )
     kda_o_cutedsl(
         qg,
@@ -219,7 +225,35 @@ def check_shape(T, H, K, V, device, dtype, num_sms):
     print(
         f"  [{status}] {tag} | o_err {o_err:.2e}  state_err {s_err:.2e}  finite={finite}"
     )
-    return ok
+
+    helion_state = torch.zeros(1, H, V, K, device=device, dtype=torch.float32)
+    helion_o = helion_chunk_kda(
+        q=inp["q"],
+        k=inp["k"],
+        v=inp["v"].clone(),
+        g=inp["g_act"],
+        beta=inp["beta"],
+        scale=scale,
+        initial_state=helion_state,
+        initial_state_indices=torch.zeros(1, device=device, dtype=torch.int32),
+        use_qk_l2norm_in_kernel=False,
+        cu_seqlens=None,
+        A_log=None,
+        dt_bias=None,
+        lower_bound=None,
+    )
+    helion_finite = bool(
+        torch.isfinite(helion_o).all() and torch.isfinite(helion_state).all()
+    )
+    helion_o_err = (helion_o[0].float() - o_ref.float()).abs().max().item()
+    helion_s_err = (helion_state.float() - state_ref.float()).abs().max().item()
+    helion_ok = helion_finite and helion_o_err < 1e-2 and helion_s_err < 5e-2
+    print(
+        f"  [{'PASS' if helion_ok else 'FAIL'}] Helion {tag} | "
+        f"o_err {helion_o_err:.2e}  state_err {helion_s_err:.2e}  "
+        f"finite={helion_finite}"
+    )
+    return ok and helion_ok
 
 
 # ---------------------------------------------------------------------------
@@ -264,15 +298,37 @@ def bench_shape(T, H, K, V, device, dtype, num_sms):
     def fn_cutedsl():
         run_cutedsl_pipeline(inp, buf, scale)
 
+    helion_v = v.clone()
+    helion_state = torch.zeros(1, H, V, K, device=device, dtype=torch.float32)
+
+    def fn_helion():
+        helion_chunk_kda(
+            q=q,
+            k=k,
+            v=helion_v,
+            g=g_act,
+            beta=beta,
+            scale=scale,
+            initial_state=helion_state,
+            initial_state_indices=idx,
+            use_qk_l2norm_in_kernel=False,
+            cu_seqlens=None,
+            A_log=None,
+            dt_bias=None,
+            lower_bound=None,
+        )
+
     quantiles = [0.5, 0.2, 0.8]
     fn_triton()
     fn_cutedsl()
+    fn_helion()
     torch.cuda.synchronize()
 
     ms_triton, _, _ = triton.testing.do_bench_cudagraph(fn_triton, quantiles=quantiles)
     ms_cutedsl, _, _ = triton.testing.do_bench_cudagraph(
         fn_cutedsl, quantiles=quantiles
     )
+    ms_helion, _, _ = triton.testing.do_bench_cudagraph(fn_helion, quantiles=quantiles)
 
     flops = kda_flops(T, H, K, V)
     mem_bytes = kda_bytes(T, H, K, V, 1, dtype)
@@ -281,7 +337,10 @@ def bench_shape(T, H, K, V, device, dtype, num_sms):
         f"  {H:>3}  {T:>7} | "
         f"{ms_triton:>8.3f}  {flops / ms_triton / 1e9:>7.2f}  {mem_bytes / ms_triton / 1e9:>7.2f} | "
         f"{ms_cutedsl:>8.3f}  {flops / ms_cutedsl / 1e9:>7.2f}  {mem_bytes / ms_cutedsl / 1e9:>7.2f} | "
-        f"{speedup:>7.2f}x"
+        f"{speedup:>7.2f}x | "
+        f"{ms_helion:>8.3f}  {flops / ms_helion / 1e9:>7.2f}  "
+        f"{mem_bytes / ms_helion / 1e9:>7.2f} | "
+        f"{ms_triton / ms_helion:>7.2f}x"
     )
 
 
@@ -312,8 +371,9 @@ def run_benchmark(device, dtype, args, num_sms):
         f"  {'H':>3}  {'T':>7} | "
         f"{'tri(ms)':>8}  {'TFLOP':>7}  {'TB/s':>7} | "
         f"{'cute(ms)':>8}  {'TFLOP':>7}  {'TB/s':>7} | {'speedup':>8}"
+        f" | {'helion(ms)':>10}  {'TFLOP':>7}  {'TB/s':>7} | {'tri/hel':>8}"
     )
-    print("  " + "-" * 84)
+    print("  " + "-" * 126)
     for H in args.num_heads:
         for T in args.seq_lens:
             bench_shape(T, H, 128, 128, device, dtype, num_sms)

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

@@ -0,0 +1,5 @@
+"""Helion attention kernels."""
+
+# K3 exposes 12 local value heads at TP=8. Lower value-head counts share the
+# same small-head decode regime.
+KDA_SMALL_VALUE_HEAD_THRESHOLD = 12

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

@@ -0,0 +1,384 @@
+"""Helion implementation of SGLang's packed KDA decode contract."""
+
+from __future__ import annotations
+
+import helion
+import helion.language as hl
+import torch
+
+from sglang.kernels.ops.attention.helion import KDA_SMALL_VALUE_HEAD_THRESHOLD
+
+# SGLang initializes torch.distributed, but this kernel has no collectives.
+_IGNORED_WARNINGS = [helion.exc.ProcessGroupNameNotFound]
+_LOG2_E = 1.4426950408889634
+
+# Tile V on the CUDA x axis so tensor-parallel head counts do not change the
+# grid width.
+_KDA_CONFIG = helion.Config(
+    block_sizes=[8],
+    loop_orders=[[2, 1, 0]],
+    num_warps=1,
+    num_stages=1,
+    indexing="pointer",
+    pid_type="xyz",
+)
+
+_KDA_BF16_CONFIG = helion.Config(
+    atomic_indexing=[],
+    block_sizes=[16],
+    indexing="pointer",
+    l2_groupings=[16],
+    # Policies are positional in the traced load order. Retune them if the
+    # decode body gains, loses, or reorders loads.
+    load_eviction_policies=[
+        "",
+        "last",
+        "first",
+        "last",
+        "first",
+        "first",
+        "last",
+        "",
+        "last",
+    ],
+    loop_orders=[[1, 2, 0]],
+    num_stages=1,
+    num_warps=1,
+    pid_type="flat",
+    range_flattens=[None],
+    range_multi_buffers=[None],
+    range_num_stages=[],
+    range_unroll_factors=[0],
+)
+
+# The bounded sigmoid gate has lower ALU and register pressure than the
+# unbounded softplus gate, allowing small-head BF16 decode to use a wider V
+# tile. The same tile regresses the unbounded path.
+_KDA_BF16_SMALL_HEAD_CONFIG = helion.Config(
+    block_sizes=[32],
+    loop_orders=[[2, 1, 0]],
+    num_warps=1,
+    num_stages=1,
+    indexing="pointer",
+    pid_type="xyz",
+)
+
+
+def _helion_fused_recurrent_kda_packed_decode_body(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    scale: float,
+    lower_bound: float,
+    initial_state: torch.Tensor,
+    out: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+    use_qk_l2norm_in_kernel: hl.constexpr = False,  # pyrefly: ignore[bad-function-definition]
+    use_fast_rsqrt: hl.constexpr = False,  # pyrefly: ignore[bad-function-definition]
+    use_lower_bound: hl.constexpr = False,  # pyrefly: ignore[bad-function-definition]
+) -> torch.Tensor:
+    """Fused packed KDA decode body; mutates ``initial_state`` and ``out``."""
+    B = mixed_qkv.size(0)
+    HV = hl.specialize(initial_state.size(-3))
+    V = hl.specialize(initial_state.size(-2))
+    K = hl.specialize(initial_state.size(-1))
+    H = hl.specialize((mixed_qkv.size(1) - HV * V) // (2 * K))
+    heads_per_q = HV // H
+
+    hl.specialize(
+        (
+            mixed_qkv.stride(0),
+            mixed_qkv.stride(1),
+            a.stride(0),
+            a.stride(1),
+            b.stride(0),
+            b.stride(1),
+            A_log.stride(0),
+            dt_bias.stride(0),
+            initial_state.stride(0),
+            initial_state.stride(1),
+            initial_state.stride(2),
+            initial_state.stride(3),
+            out.stride(0),
+            out.stride(1),
+            out.stride(2),
+            out.stride(3),
+            ssm_state_indices.stride(0),
+        )
+    )
+
+    block_v = hl.register_block_size(1, V)
+
+    for tile_b, tile_hv, tile_v in hl.tile([B, HV, V], block_size=[1, 1, block_v]):
+        k_offsets = hl.arange(K)
+        i_b = tile_b.id
+        i_hv = tile_hv.id
+        i_h = i_hv // heads_per_q
+
+        state_index = ssm_state_indices[i_b].long()
+        if state_index < 0:
+            out[i_b, 0, i_hv, tile_v] = 0.0
+        else:
+            q_offsets = i_h * K + k_offsets
+            k_input_offsets = H * K + i_h * K + k_offsets
+            v_offsets = 2 * H * K + i_hv * V + tile_v.index
+
+            raw_gate = a[i_b, i_hv * K + k_offsets].float()
+            raw_gate = raw_gate + dt_bias[i_hv * K + k_offsets].float()
+            A_log_value = A_log[i_hv].float()
+            A = torch.exp2(A_log_value * _LOG2_E)
+            if use_lower_bound:
+                log_decay = lower_bound * torch.sigmoid(A * raw_gate)
+            else:
+                gate_exp = torch.exp2(raw_gate * _LOG2_E)
+                softplus = torch.where(
+                    raw_gate <= 20.0,
+                    torch.log(1.0 + gate_exp),
+                    raw_gate,
+                )
+                log_decay = -A * softplus
+            beta = torch.sigmoid(b[i_b, i_hv].float())
+
+            state = initial_state[state_index, i_hv, tile_v.index, k_offsets].float()
+            decay = torch.exp2(log_decay * _LOG2_E)
+            state = state * decay[None, :]
+
+            k = mixed_qkv[i_b, k_input_offsets].float()
+            if use_qk_l2norm_in_kernel:
+                k_norm = (k * k).sum() + 1e-6
+                if use_fast_rsqrt:
+                    k = k * torch.rsqrt(k_norm)
+                else:
+                    k = k / torch.sqrt(k_norm)
+            v = mixed_qkv[i_b, v_offsets].float()
+            value_residual = v - (state * k[None, :]).sum(-1)
+            value_residual = value_residual * beta
+            state = state + value_residual[:, None] * k[None, :]
+
+            q = mixed_qkv[i_b, q_offsets].float()
+            if use_qk_l2norm_in_kernel:
+                q_norm = (q * q).sum() + 1e-6
+                if use_fast_rsqrt:
+                    q = q * torch.rsqrt(q_norm)
+                else:
+                    q = q / torch.sqrt(q_norm)
+            q = q * scale
+            output = (state * q[None, :]).sum(-1)
+
+            out[i_b, 0, i_hv, tile_v] = output.to(out.dtype)
+            initial_state[state_index, i_hv, tile_v.index, k_offsets] = state
+
+    return out
+
+
+_helion_fused_recurrent_kda_packed_decode = helion.kernel(
+    _helion_fused_recurrent_kda_packed_decode_body,
+    static_shapes=False,
+    config=_KDA_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)
+_helion_fused_recurrent_kda_packed_decode_bf16 = helion.kernel(
+    _helion_fused_recurrent_kda_packed_decode_body,
+    static_shapes=False,
+    config=_KDA_BF16_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)
+_helion_fused_recurrent_kda_packed_decode_bf16_small_head = helion.kernel(
+    _helion_fused_recurrent_kda_packed_decode_body,
+    static_shapes=False,
+    config=_KDA_BF16_SMALL_HEAD_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)
+
+
+def _select_decode_kernel(
+    *,
+    is_bf16_state: bool,
+    num_v_heads: int,
+    use_lower_bound: bool,
+) -> helion.Kernel:
+    if (
+        is_bf16_state
+        and use_lower_bound
+        and num_v_heads <= KDA_SMALL_VALUE_HEAD_THRESHOLD
+    ):
+        return _helion_fused_recurrent_kda_packed_decode_bf16_small_head
+    if is_bf16_state:
+        return _helion_fused_recurrent_kda_packed_decode_bf16
+    return _helion_fused_recurrent_kda_packed_decode
+
+
+def validate_packed_decode_inputs(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    initial_state: torch.Tensor,
+    out: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+) -> tuple[int, int, int, int, int]:
+    """Apply the shape and layout checks from SGLang's packed wrapper."""
+    if mixed_qkv.ndim != 2:
+        raise ValueError(
+            f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
+        )
+    if mixed_qkv.stride(-1) != 1:
+        raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
+    if a.ndim != 2 or b.ndim != 2:
+        raise ValueError(
+            f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
+        )
+    if a.stride(-1) != 1 or b.stride(-1) != 1:
+        raise ValueError("`a`/`b` must be contiguous in the last dim.")
+    if A_log.ndim != 1 or dt_bias.ndim != 1:
+        raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
+    if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
+        raise ValueError("`A_log`/`dt_bias` must be contiguous.")
+    if ssm_state_indices.ndim != 1:
+        raise ValueError(
+            "`ssm_state_indices` must be 1D for packed decode "
+            f"(got ndim={ssm_state_indices.ndim})."
+        )
+    if not out.is_contiguous():
+        raise ValueError("`out` must be contiguous.")
+
+    device = mixed_qkv.device
+    if any(
+        tensor.device != device
+        for tensor in (
+            a,
+            b,
+            A_log,
+            dt_bias,
+            initial_state,
+            out,
+            ssm_state_indices,
+        )
+    ):
+        raise ValueError("All inputs must be on the same device.")
+
+    B = mixed_qkv.shape[0]
+    if a.shape[0] != B or b.shape[0] != B:
+        raise ValueError(
+            "Mismatched batch sizes: "
+            f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, "
+            f"b.shape[0]={b.shape[0]}."
+        )
+    if ssm_state_indices.shape[0] != B:
+        raise ValueError(
+            f"`ssm_state_indices` must have shape [B] "
+            f"(got {tuple(ssm_state_indices.shape)}; expected ({B},))."
+        )
+
+    if initial_state.ndim != 4:
+        raise ValueError(
+            f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
+        )
+    if initial_state.stride(-1) != 1:
+        raise ValueError("`initial_state` must be contiguous in the last dim.")
+    HV, V, K = initial_state.shape[-3:]
+    if a.shape[1] != HV * K:
+        raise ValueError(
+            f"`a` must have shape [B, HV*K] with HV={HV}, K={K} "
+            f"(got a.shape={tuple(a.shape)})."
+        )
+    if b.shape[1] != HV:
+        raise ValueError(
+            f"`b` must have shape [B, HV] with HV={HV} (got b.shape={tuple(b.shape)})."
+        )
+    if A_log.numel() != HV:
+        raise ValueError(f"`A_log` must have {HV} elements (got {A_log.numel()}).")
+    if dt_bias.numel() != HV * K:
+        raise ValueError(
+            f"`dt_bias` must have {HV * K} elements (got {dt_bias.numel()})."
+        )
+    if out.shape != (B, 1, HV, V):
+        raise ValueError(
+            f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})."
+        )
+
+    qkv_dim = mixed_qkv.shape[1]
+    qk_dim = qkv_dim - HV * V
+    if qk_dim <= 0 or qk_dim % 2 != 0:
+        raise ValueError(
+            f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}."
+        )
+    q_dim = qk_dim // 2
+    if q_dim % K != 0:
+        raise ValueError(
+            f"Invalid packed Q size {q_dim}: must be divisible by K={K}. "
+            "KDA packed decode requires num_q_heads == num_k_heads and "
+            "head_q_dim == head_k_dim."
+        )
+    H = q_dim // K
+    if H <= 0 or HV % H != 0:
+        raise ValueError(
+            f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}."
+        )
+    return B, H, HV, K, V
+
+
+def helion_fused_recurrent_kda_packed_decode(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    out: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+    use_qk_l2norm_in_kernel: bool = False,
+    lower_bound: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Helion implementation of SGLang's packed KDA decode contract.
+
+    Inputs, mutations, padding semantics, and outputs match
+    ``fused_recurrent_kda_packed_decode``:
+
+    * ``mixed_qkv`` is ``[B, 2*H*K + HV*V]`` after the short convolution.
+    * ``a`` and ``b`` are raw forget-gate and beta logits.
+    * ``lower_bound`` selects the bounded sigmoid decay used by safe-gate KDA.
+    * ``initial_state`` is ``[num_slots, HV, V, K]`` and is updated in place.
+    * ``ssm_state_indices == -1`` writes a zero output and leaves state untouched.
+    * ``out`` is ``[B, 1, HV, V]`` and is written in place.
+    * The return is the same ``(out, initial_state)`` object pair supplied by the
+      caller.
+    """
+    _, _, num_v_heads, _, _ = validate_packed_decode_inputs(
+        mixed_qkv,
+        a,
+        b,
+        A_log,
+        dt_bias,
+        initial_state,
+        out,
+        ssm_state_indices,
+    )
+    use_lower_bound = lower_bound is not None
+    is_bf16_state = initial_state.dtype is torch.bfloat16
+    kernel = _select_decode_kernel(
+        is_bf16_state=is_bf16_state,
+        num_v_heads=num_v_heads,
+        use_lower_bound=use_lower_bound,
+    )
+    lower_bound_value = 0.0 if lower_bound is None else lower_bound
+    result = kernel(
+        mixed_qkv,
+        a,
+        b,
+        A_log,
+        dt_bias,
+        scale,
+        lower_bound_value,
+        initial_state,
+        out,
+        ssm_state_indices,
+        use_qk_l2norm_in_kernel,
+        is_bf16_state,
+        use_lower_bound,
+    )
+    return result, initial_state

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

@@ -0,0 +1,780 @@
+"""Helion ReplaySSM decode for Kimi Delta Attention.
+
+The public :func:`helion_fused_recurrent_kda_replayssm_decode` mirrors
+``fused_recurrent_linear_replayssm_decode(..., is_kda=True)`` from
+``sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm``.  That
+Triton kernel is gate-generic (GDN scalar gate / KDA per-K gate); this module
+implements the KDA path only, and additionally supports the bounded gate,
+which the Triton ReplaySSM kernel does not expose.
+"""
+
+from __future__ import annotations
+
+import helion
+import helion.language as hl
+import torch
+
+from sglang.kernels.ops.attention.helion import KDA_SMALL_VALUE_HEAD_THRESHOLD
+from sglang.kernels.ops.attention.helion.kda_decode import (
+    validate_packed_decode_inputs,
+)
+
+# SGLang initializes torch.distributed, but this kernel has no collectives.
+_IGNORED_WARNINGS = [helion.exc.ProcessGroupNameNotFound]
+_LOG2_E = 1.4426950408889634
+
+# Indexing and eviction-policy lists are positional in the traced load order.
+# Retune them if the ReplaySSM body gains, loses, or reorders loads.
+_KDA_REPLAYSSM_FP32_CONFIG = helion.Config(
+    atomic_indexing=[],
+    block_sizes=[32, 128],
+    indexing=[
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "pointer",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "pointer",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "tensor_descriptor",
+    ],
+    l2_groupings=[4],
+    load_eviction_policies=[
+        "last",
+        "last",
+        "first",
+        "first",
+        "",
+        "first",
+        "",
+        "first",
+        "first",
+        "",
+        "last",
+        "first",
+        "last",
+        "first",
+        "last",
+        "first",
+        "",
+        "first",
+        "last",
+        "",
+        "last",
+        "first",
+        "last",
+        "first",
+        "first",
+    ],
+    loop_orders=[[1, 2, 0]],
+    num_warps=1,
+    num_stages=1,
+    pid_type="flat",
+    range_flattens=[None, None],
+    range_multi_buffers=[None, False],
+    range_num_stages=[0, 3],
+    range_unroll_factors=[0, 0],
+)
+
+_KDA_REPLAYSSM_BF16_CONFIG = helion.Config(
+    atomic_indexing=[],
+    block_sizes=[64, 64],
+    indexing=[
+        "pointer",
+        "pointer",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "pointer",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "tensor_descriptor",
+        "pointer",
+        "tensor_descriptor",
+        "pointer",
+        "pointer",
+        "pointer",
+        "pointer",
+        "pointer",
+        "pointer",
+        "pointer",
+        "pointer",
+    ],
+    l2_groupings=[2],
+    load_eviction_policies=[
+        "first",
+        "last",
+        "last",
+        "first",
+        "",
+        "",
+        "last",
+        "last",
+        "first",
+        "first",
+        "first",
+        "",
+        "last",
+        "first",
+        "last",
+        "first",
+        "",
+        "first",
+        "",
+        "first",
+        "first",
+        "",
+        "",
+        "first",
+        "last",
+    ],
+    loop_orders=[[1, 2, 0]],
+    num_stages=1,
+    num_warps=1,
+    pid_type="flat",
+    range_flattens=[None, False],
+    range_multi_buffers=[None, False],
+    range_num_stages=[0, 0],
+    range_unroll_factors=[0, 4],
+)
+
+# Small-head BF16 ReplaySSM uses the FP32 tile schedule with direct PID order.
+_KDA_REPLAYSSM_BF16_SMALL_HEAD_CONFIG = helion.Config.from_dict(
+    {**_KDA_REPLAYSSM_FP32_CONFIG, "l2_groupings": [1]}
+)
+
+
+def _log_decay(
+    *,
+    raw_gate: torch.Tensor,
+    decay_rate: torch.Tensor,
+    lower_bound: float,
+    use_lower_bound: hl.constexpr,
+) -> torch.Tensor:
+    if use_lower_bound:
+        return lower_bound * torch.sigmoid(decay_rate * raw_gate)
+    gate_exp = torch.exp2(raw_gate * _LOG2_E)
+    softplus = torch.where(
+        raw_gate <= 20.0,
+        torch.log(1.0 + gate_exp),
+        raw_gate,
+    )
+    return -decay_rate * softplus
+
+
+def _load_log_decay(
+    *,
+    a: torch.Tensor,
+    dt_bias: torch.Tensor,
+    batch_index: torch.Tensor,
+    value_head: torch.Tensor,
+    key_indices: torch.Tensor,
+    key_valid: torch.Tensor,
+    key_dim: int,
+    decay_rate: torch.Tensor,
+    lower_bound: float,
+    use_lower_bound: hl.constexpr,
+) -> torch.Tensor:
+    """Load the current raw gate and apply the selected KDA gate contract."""
+    raw_gate = hl.load(
+        a,
+        [batch_index, value_head * key_dim + key_indices],
+        extra_mask=key_valid,
+    ).float()
+    raw_gate = (
+        raw_gate
+        + hl.load(
+            dt_bias,
+            [value_head * key_dim + key_indices],
+            extra_mask=key_valid,
+        ).float()
+    )
+    return _log_decay(
+        raw_gate=raw_gate,
+        decay_rate=decay_rate,
+        lower_bound=lower_bound,
+        use_lower_bound=use_lower_bound,
+    )
+
+
+def _helion_fused_recurrent_kda_replayssm_decode_body(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    d_cache: torch.Tensor,
+    k_cache: torch.Tensor,
+    g_cache: torch.Tensor,
+    out: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+    write_pos: torch.Tensor,
+    force_flush: torch.Tensor | None,
+    lower_bound: float,
+    cache_block: hl.constexpr,  # pyrefly: ignore[bad-function-definition]
+    use_qk_l2norm_in_kernel: hl.constexpr,  # pyrefly: ignore[bad-function-definition]
+    use_lower_bound: hl.constexpr,  # pyrefly: ignore[bad-function-definition]
+) -> torch.Tensor:
+    """Reconstruct the buffered state, emit one token, and flush if needed."""
+    B = mixed_qkv.size(0)
+    HV = hl.specialize(initial_state.size(1))
+    V = hl.specialize(initial_state.size(2))
+    K = hl.specialize(initial_state.size(3))
+    H = hl.specialize(k_cache.size(1))
+    cache_length = hl.specialize(d_cache.size(2))
+    heads_per_q = HV // H
+
+    # Keep storage setup inline so the generated host wrapper does not call
+    # Python helpers on every eager invocation.
+    state_strides = hl.specialize(
+        (
+            initial_state.stride(0),
+            initial_state.stride(1),
+            initial_state.stride(2),
+            initial_state.stride(3),
+        )
+    )
+    d_strides = hl.specialize(
+        (
+            d_cache.stride(0),
+            d_cache.stride(1),
+            d_cache.stride(2),
+            d_cache.stride(3),
+        )
+    )
+    k_strides = hl.specialize(
+        (
+            k_cache.stride(0),
+            k_cache.stride(1),
+            k_cache.stride(2),
+            k_cache.stride(3),
+        )
+    )
+    g_strides = hl.specialize(
+        (
+            g_cache.stride(0),
+            g_cache.stride(1),
+            g_cache.stride(2),
+            g_cache.stride(3),
+        )
+    )
+    state_storage_size = (
+        (initial_state.size(0) - 1) * state_strides[0]
+        + (HV - 1) * state_strides[1]
+        + (V - 1) * state_strides[2]
+        + (K - 1) * state_strides[3]
+        + 1
+    )
+    d_storage_size = (
+        (d_cache.size(0) - 1) * d_strides[0]
+        + (HV - 1) * d_strides[1]
+        + (cache_length - 1) * d_strides[2]
+        + (V - 1) * d_strides[3]
+        + 1
+    )
+    k_storage_size = (
+        (k_cache.size(0) - 1) * k_strides[0]
+        + (H - 1) * k_strides[1]
+        + (cache_length - 1) * k_strides[2]
+        + (K - 1) * k_strides[3]
+        + 1
+    )
+    g_storage_size = (
+        (g_cache.size(0) - 1) * g_strides[0]
+        + (HV - 1) * g_strides[1]
+        + (cache_length - 1) * g_strides[2]
+        + (K - 1) * g_strides[3]
+        + 1
+    )
+    # Omitting storage_offset preserves each source view's existing offset,
+    # which is required for envelope and page-major cache views.
+    state_storage = initial_state.as_strided([state_storage_size], [1])
+    d_storage = d_cache.as_strided([d_storage_size], [1])
+    k_storage = k_cache.as_strided([k_storage_size], [1])
+    g_storage = g_cache.as_strided([g_storage_size], [1])
+
+    hl.specialize(
+        (
+            mixed_qkv.stride(0),
+            mixed_qkv.stride(1),
+            a.stride(0),
+            a.stride(1),
+            b.stride(0),
+            b.stride(1),
+            A_log.stride(0),
+            dt_bias.stride(0),
+            out.stride(0),
+            out.stride(1),
+            out.stride(2),
+            out.stride(3),
+            ssm_state_indices.stride(0),
+            write_pos.stride(0),
+        )
+    )
+
+    block_v = hl.register_block_size(16, V)
+    block_k = hl.register_block_size(16, K)
+
+    for tile_b, tile_hv, tile_v in hl.tile([B, HV, V], block_size=[1, 1, block_v]):
+        i_b = tile_b.id
+        i_hv = tile_hv.id
+        i_h = i_hv // heads_per_q
+        state_index = ssm_state_indices[i_b].long()
+        v_valid = tile_v.index < V
+
+        if state_index < 0:
+            hl.store(
+                out,
+                [i_b, 0, i_hv, tile_v.index],
+                hl.zeros([tile_v], dtype=out.dtype),
+                extra_mask=v_valid,
+            )
+        else:
+            cursor = write_pos[i_b].long()
+            is_flush = cursor == cache_length - 1
+            if force_flush is not None:
+                is_flush = is_flush | (force_flush[i_b] != 0)
+            should_append = is_flush == 0
+
+            cache_positions = hl.arange(cache_block)
+            cache_valid = cache_positions < cursor
+            d_offsets = (
+                state_index * d_strides[0]
+                + i_hv * d_strides[1]
+                + cache_positions[:, None] * d_strides[2]
+                + tile_v.index[None, :] * d_strides[3]
+            )
+            d_values = hl.load(
+                d_storage,
+                [d_offsets],
+                extra_mask=cache_valid[:, None] & v_valid[None, :],
+            ).float()
+            # Advanced loads preserve tensor-axis order: d_cache contributes
+            # [L, V], while the reconstruction dot consumes [V, L].
+            d_dot = d_values.T.to(out.dtype)
+
+            full_k = hl.arange(K)
+            q_offsets = i_h * K + full_k
+            full_k_offsets = H * K + i_h * K + full_k
+            q_full = mixed_qkv[i_b, q_offsets].float()
+            k_full = mixed_qkv[i_b, full_k_offsets].float()
+            if use_qk_l2norm_in_kernel:
+                q_rnorm = 1.0 / torch.sqrt((q_full * q_full).sum() + 1e-6)
+                k_rnorm = 1.0 / torch.sqrt((k_full * k_full).sum() + 1e-6)
+            else:
+                q_rnorm = 1.0
+                k_rnorm = 1.0
+
+            value_offsets = 2 * H * K + i_hv * V + tile_v.index
+            value = hl.load(
+                mixed_qkv,
+                [i_b, value_offsets],
+                extra_mask=v_valid,
+            ).float()
+            # ReplaySSM Triton rounds beta through the input dtype before FP32
+            # accumulation; packed KDA intentionally keeps beta in FP32.
+            beta = torch.sigmoid(b[i_b, i_hv].float()).to(b.dtype).float()
+            A = torch.exp2(A_log[i_hv].float() * _LOG2_E)
+
+            state_q = hl.zeros([tile_v], dtype=torch.float32)
+            state_k = hl.zeros([tile_v], dtype=torch.float32)
+            current_kq = hl.zeros([], dtype=torch.float32)
+
+            # Reconstruct each K tile directly from the checkpoint and ring.
+            # The full state is never materialized outside registers.
+            for tile_k in hl.tile(K, block_size=block_k):
+                k_valid = tile_k.index < K
+                q_value = hl.load(
+                    mixed_qkv,
+                    [i_b, i_h * K + tile_k.index],
+                    extra_mask=k_valid,
+                ).float()
+                k_value = hl.load(
+                    mixed_qkv,
+                    [i_b, H * K + i_h * K + tile_k.index],
+                    extra_mask=k_valid,
+                ).float()
+                q_value = q_value * q_rnorm
+                k_value = k_value * k_rnorm
+                q_scaled = q_value * scale
+                current_kq = current_kq + (k_value * q_scaled).sum()
+
+                cache_mask = cache_valid[:, None] & k_valid[None, :]
+                g_offsets = (
+                    state_index * g_strides[0]
+                    + i_hv * g_strides[1]
+                    + cache_positions[:, None] * g_strides[2]
+                    + tile_k.index[None, :] * g_strides[3]
+                )
+                cached_k_offsets = (
+                    state_index * k_strides[0]
+                    + i_h * k_strides[1]
+                    + cache_positions[:, None] * k_strides[2]
+                    + tile_k.index[None, :] * k_strides[3]
+                )
+                state_offsets = (
+                    state_index * state_strides[0]
+                    + i_hv * state_strides[1]
+                    + tile_v.index[:, None] * state_strides[2]
+                    + tile_k.index[None, :] * state_strides[3]
+                )
+                state_mask = v_valid[:, None] & k_valid[None, :]
+                cached_g = hl.load(
+                    g_storage,
+                    [g_offsets],
+                    extra_mask=cache_mask,
+                ).float()
+                gate_prefix = torch.cumsum(cached_g, dim=0)
+                gate_total = cached_g.sum(0)
+                replay_decay = torch.where(
+                    cache_mask,
+                    torch.exp2((gate_total[None, :] - gate_prefix) * _LOG2_E),
+                    0.0,
+                )
+                total_decay = torch.exp2(gate_total * _LOG2_E)
+                cached_k = hl.load(
+                    k_storage,
+                    [cached_k_offsets],
+                    extra_mask=cache_mask,
+                ).float()
+                cached_k = (cached_k * replay_decay).to(out.dtype)
+                state = hl.load(
+                    state_storage,
+                    [state_offsets],
+                    extra_mask=state_mask,
+                ).float()
+                state = state * total_decay[None, :]
+                state = state + hl.dot(
+                    d_dot,
+                    cached_k,
+                    out_dtype=torch.float32,
+                )
+
+                current_gate = _load_log_decay(
+                    a=a,
+                    dt_bias=dt_bias,
+                    batch_index=i_b,
+                    value_head=i_hv,
+                    key_indices=tile_k.index,
+                    key_valid=k_valid,
+                    key_dim=K,
+                    decay_rate=A,
+                    lower_bound=lower_bound,
+                    use_lower_bound=use_lower_bound,
+                )
+                current_decay = torch.exp2(current_gate * _LOG2_E)
+                q_effective = q_scaled * current_decay
+                k_effective = k_value * current_decay
+                state_q = state_q + (state * q_effective[None, :]).sum(-1)
+                state_k = state_k + (state * k_effective[None, :]).sum(-1)
+
+                if should_append:
+                    if tile_v.id == 0:
+                        if i_hv == i_h * heads_per_q:
+                            current_k_offsets = (
+                                state_index * k_strides[0]
+                                + i_h * k_strides[1]
+                                + cursor * k_strides[2]
+                                + tile_k.index * k_strides[3]
+                            )
+                            hl.store(
+                                k_storage,
+                                [current_k_offsets],
+                                k_value.to(k_cache.dtype),
+                                extra_mask=(cursor < cache_length) & k_valid,
+                            )
+                        current_g_offsets = (
+                            state_index * g_strides[0]
+                            + i_hv * g_strides[1]
+                            + cursor * g_strides[2]
+                            + tile_k.index * g_strides[3]
+                        )
+                        hl.store(
+                            g_storage,
+                            [current_g_offsets],
+                            current_gate,
+                            extra_mask=(cursor < cache_length) & k_valid,
+                        )
+
+            delta = beta * (value - state_k)
+            output = state_q + delta * current_kq
+            hl.store(
+                out,
+                [i_b, 0, i_hv, tile_v.index],
+                output.to(out.dtype),
+                extra_mask=v_valid,
+            )
+
+            if is_flush:
+                # Reconstruct again now that the current delta is available,
+                # then fold the current rank-one update into the checkpoint.
+                # Helion cannot lower the cumsum through a device helper without
+                # creating a separate unsupported scan subgraph.
+                for tile_k in hl.tile(K, block_size=block_k):
+                    k_valid = tile_k.index < K
+                    k_value = hl.load(
+                        mixed_qkv,
+                        [i_b, H * K + i_h * K + tile_k.index],
+                        extra_mask=k_valid,
+                    ).float()
+                    k_value = k_value * k_rnorm
+
+                    cache_mask = cache_valid[:, None] & k_valid[None, :]
+                    g_offsets = (
+                        state_index * g_strides[0]
+                        + i_hv * g_strides[1]
+                        + cache_positions[:, None] * g_strides[2]
+                        + tile_k.index[None, :] * g_strides[3]
+                    )
+                    cached_k_offsets = (
+                        state_index * k_strides[0]
+                        + i_h * k_strides[1]
+                        + cache_positions[:, None] * k_strides[2]
+                        + tile_k.index[None, :] * k_strides[3]
+                    )
+                    state_offsets = (
+                        state_index * state_strides[0]
+                        + i_hv * state_strides[1]
+                        + tile_v.index[:, None] * state_strides[2]
+                        + tile_k.index[None, :] * state_strides[3]
+                    )
+                    state_mask = v_valid[:, None] & k_valid[None, :]
+                    cached_g = hl.load(
+                        g_storage,
+                        [g_offsets],
+                        extra_mask=cache_mask,
+                    ).float()
+                    gate_prefix = torch.cumsum(cached_g, dim=0)
+                    gate_total = cached_g.sum(0)
+                    replay_decay = torch.where(
+                        cache_mask,
+                        torch.exp2((gate_total[None, :] - gate_prefix) * _LOG2_E),
+                        0.0,
+                    )
+                    total_decay = torch.exp2(gate_total * _LOG2_E)
+                    cached_k = hl.load(
+                        k_storage,
+                        [cached_k_offsets],
+                        extra_mask=cache_mask,
+                    ).float()
+                    cached_k = (cached_k * replay_decay).to(out.dtype)
+                    state = hl.load(
+                        state_storage,
+                        [state_offsets],
+                        extra_mask=state_mask,
+                    ).float()
+                    state = state * total_decay[None, :]
+                    state = state + hl.dot(
+                        d_dot,
+                        cached_k,
+                        out_dtype=torch.float32,
+                    )
+                    current_gate = _load_log_decay(
+                        a=a,
+                        dt_bias=dt_bias,
+                        batch_index=i_b,
+                        value_head=i_hv,
+                        key_indices=tile_k.index,
+                        key_valid=k_valid,
+                        key_dim=K,
+                        decay_rate=A,
+                        lower_bound=lower_bound,
+                        use_lower_bound=use_lower_bound,
+                    )
+                    current_decay = torch.exp2(current_gate * _LOG2_E)
+                    state = state * current_decay[None, :]
+                    state = state + delta[:, None] * k_value[None, :]
+                    hl.store(
+                        state_storage,
+                        [state_offsets],
+                        state.to(initial_state.dtype),
+                        extra_mask=v_valid[:, None] & k_valid[None, :],
+                    )
+            else:
+                current_d_offsets = (
+                    state_index * d_strides[0]
+                    + i_hv * d_strides[1]
+                    + cursor * d_strides[2]
+                    + tile_v.index * d_strides[3]
+                )
+                hl.store(
+                    d_storage,
+                    [current_d_offsets],
+                    delta.to(d_cache.dtype),
+                    extra_mask=(cursor < cache_length) & v_valid,
+                )
+
+    return out
+
+
+_helion_fused_recurrent_kda_replayssm_decode_fp32 = helion.kernel(
+    _helion_fused_recurrent_kda_replayssm_decode_body,
+    static_shapes=False,
+    config=_KDA_REPLAYSSM_FP32_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)
+_helion_fused_recurrent_kda_replayssm_decode_bf16 = helion.kernel(
+    _helion_fused_recurrent_kda_replayssm_decode_body,
+    static_shapes=False,
+    config=_KDA_REPLAYSSM_BF16_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)
+_helion_fused_recurrent_kda_replayssm_decode_bf16_small_head = helion.kernel(
+    _helion_fused_recurrent_kda_replayssm_decode_body,
+    static_shapes=False,
+    config=_KDA_REPLAYSSM_BF16_SMALL_HEAD_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)
+
+
+def _select_replayssm_decode_kernel(
+    *,
+    is_bf16_state: bool,
+    num_v_heads: int,
+) -> helion.Kernel:
+    if is_bf16_state and num_v_heads <= KDA_SMALL_VALUE_HEAD_THRESHOLD:
+        return _helion_fused_recurrent_kda_replayssm_decode_bf16_small_head
+    if is_bf16_state:
+        return _helion_fused_recurrent_kda_replayssm_decode_bf16
+    return _helion_fused_recurrent_kda_replayssm_decode_fp32
+
+
+def helion_fused_recurrent_kda_replayssm_decode(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    d_cache: torch.Tensor,
+    k_cache: torch.Tensor,
+    g_cache: torch.Tensor,
+    out: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+    write_pos: torch.Tensor,
+    force_flush: torch.Tensor | None = None,
+    use_qk_l2norm_in_kernel: bool = False,
+    lower_bound: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run one buffered KDA decode step using caller-owned ReplaySSM state.
+
+    Allocates nothing persistent: the caller owns ``d_cache`` / ``k_cache`` /
+    ``g_cache`` and is responsible for advancing ``write_pos`` modulo the ring
+    length after a non-flush step and resetting it to zero after a natural or
+    forced flush. ``initial_state`` is both the checkpoint read (h0) and the
+    flush-only checkpoint write (ht), in place.
+    """
+    batch = mixed_qkv.size(0)
+    if a.ndim not in (2, 3) or not a.is_contiguous():
+        raise ValueError("KDA `a` must be a contiguous 2D or 3D tensor.")
+    if dt_bias.ndim not in (1, 2) or not dt_bias.is_contiguous():
+        raise ValueError("KDA `dt_bias` must be a contiguous 1D or 2D tensor.")
+    flat_a = a.view(batch, -1)
+    flat_dt_bias = dt_bias.view(-1)
+    _, num_q_heads, num_v_heads, key_dim, value_dim = validate_packed_decode_inputs(
+        mixed_qkv,
+        flat_a,
+        b,
+        A_log,
+        flat_dt_bias,
+        initial_state,
+        out,
+        ssm_state_indices,
+    )
+
+    if write_pos.ndim != 1 or write_pos.dtype is not torch.int32:
+        raise ValueError("`write_pos` must be a 1D int32 tensor.")
+    if write_pos.shape != (batch,):
+        raise ValueError(f"`write_pos` must have shape {(batch,)}.")
+    if force_flush is not None and (
+        force_flush.ndim != 1
+        or force_flush.dtype is not torch.int32
+        or force_flush.shape != (batch,)
+    ):
+        raise ValueError("`force_flush` must be a length-B int32 tensor or None.")
+
+    cache_length = d_cache.size(2)
+    if cache_length < 1:
+        raise ValueError("ReplaySSM cache length must be at least 1.")
+    if d_cache.shape[1:] != (num_v_heads, cache_length, value_dim):
+        raise ValueError("`d_cache` must have shape [slots, HV, L, V].")
+    if k_cache.shape[1:] != (num_q_heads, cache_length, key_dim):
+        raise ValueError("`k_cache` must have shape [slots, H, L, K].")
+    if g_cache.shape[1:] != (num_v_heads, cache_length, key_dim):
+        raise ValueError("`g_cache` must have shape [slots, HV, L, K].")
+    if g_cache.dtype is not torch.float32:
+        raise ValueError("`g_cache` must have dtype torch.float32.")
+
+    device = mixed_qkv.device
+    if any(
+        tensor.device != device for tensor in (d_cache, k_cache, g_cache, write_pos)
+    ):
+        raise ValueError("ReplaySSM inputs must be on the same device.")
+    if force_flush is not None and force_flush.device != device:
+        raise ValueError("`force_flush` must be on the same device as the inputs.")
+
+    cache_block = helion.next_power_of_2(max(16, cache_length))
+    use_lower_bound = lower_bound is not None
+    kernel = _select_replayssm_decode_kernel(
+        is_bf16_state=initial_state.dtype is torch.bfloat16,
+        num_v_heads=num_v_heads,
+    )
+    result = kernel(
+        mixed_qkv,
+        flat_a,
+        b,
+        A_log,
+        flat_dt_bias,
+        scale,
+        initial_state,
+        d_cache,
+        k_cache,
+        g_cache,
+        out,
+        ssm_state_indices,
+        write_pos,
+        force_flush,
+        0.0 if lower_bound is None else lower_bound,
+        cache_block,
+        use_qk_l2norm_in_kernel,
+        use_lower_bound,
+    )
+    return result, initial_state

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

@@ -137,6 +137,10 @@ class GDNKernelDispatcher:
 
             flashinfer_kernel = FlashInferGDNKernel()
             self.decode_kernel = flashinfer_kernel
+        elif decode_backend.is_helion():
+            raise ValueError(
+                "The Helion linear-attention backend supports KDA only, not GDN."
+            )
         else:
             raise ValueError(f"Unsupported GDN decode backend: {decode_backend}")
 
@@ -176,6 +180,10 @@ class GDNKernelDispatcher:
 
                 flashinfer_kernel = FlashInferGDNKernel()
                 self.extend_kernel = flashinfer_kernel
+        elif prefill_backend.is_helion():
+            raise ValueError(
+                "The Helion linear-attention backend supports KDA only, not GDN."
+            )
         else:
             raise ValueError(f"Unsupported GDN prefill backend: {prefill_backend}")
 

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

@@ -47,9 +47,25 @@ class KDAKernelDispatcher:
     ):
         self.verify_backend = verify_backend
         triton_kernel = TritonKDAKernel()
+        helion_kernel = None
+        if decode_backend.is_helion() or prefill_backend.is_helion():
+            if not is_cuda():
+                raise ValueError("KDA Helion backend requires CUDA")
+            from sglang.srt.layers.attention.linear.kernels.kda_helion import (
+                HelionKDAKernel,
+            )
+
+            helion_kernel = HelionKDAKernel(
+                triton_fallback=triton_kernel,
+                enable_decode=decode_backend.is_helion(),
+                enable_prefill=prefill_backend.is_helion(),
+            )
 
         if decode_backend.is_triton():
             self.decode_kernel = triton_kernel
+        elif decode_backend.is_helion():
+            assert helion_kernel is not None
+            self.decode_kernel = helion_kernel
         elif decode_backend.is_cutedsl():
             if not is_cuda():
                 raise ValueError("KDA CuTe DSL backend requires CUDA")
@@ -71,7 +87,7 @@ class KDAKernelDispatcher:
         else:
             raise ValueError(
                 f"Unsupported KDA decode backend: {decode_backend}. "
-                "KDA supports 'triton', 'cutedsl', or 'flashinfer'."
+                "KDA supports 'triton', 'helion', 'cutedsl', or 'flashinfer'."
             )
 
         # target_verify kernel, selected via --linear-attn-verify-backend (defaults
@@ -110,6 +126,9 @@ class KDAKernelDispatcher:
 
         if prefill_backend.is_triton():
             self.extend_kernel = triton_kernel
+        elif prefill_backend.is_helion():
+            assert helion_kernel is not None
+            self.extend_kernel = helion_kernel
         elif prefill_backend.is_flashkda():
             from sglang.srt.layers.attention.linear.kernels.kda_flashkda import (
                 FlashKDAKernel,
@@ -166,8 +185,9 @@ class KDAKernelDispatcher:
         else:
             raise ValueError(
                 f"Unsupported KDA prefill backend: {prefill_backend}. "
-                "KDA supports 'triton', 'flashkda', 'cutedsl', 'nvidia_kda', or "
+                "KDA supports 'triton', 'helion', 'flashkda', 'cutedsl', "
-                "'ptx_kda' (cutedsl/nvidia_kda prefill need SM100, ptx_kda SM103)."
+                "'nvidia_kda', or 'ptx_kda' (cutedsl/nvidia_kda prefill need "
+                "SM100, ptx_kda SM103)."
             )
 
         self.supports_packed_decode = getattr(

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

@@ -0,0 +1,191 @@
+"""Helion backend for Kimi Delta Attention."""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
+from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
+    LinearAttnKernelBase,
+)
+
+
+class HelionKDAKernel(LinearAttnKernelBase):
+    """KDA packed decode and prefill implemented with Helion kernels.
+
+    The generic decode interface delegates to Triton, and the dispatcher routes
+    speculative target verification directly to Triton. The one-token decode
+    and ReplaySSM paths use :meth:`packed_decode`, while prefill uses
+    :meth:`extend`.
+    """
+
+    supports_packed_decode = True
+
+    def __init__(
+        self,
+        triton_fallback: TritonKDAKernel | None = None,
+        *,
+        enable_decode: bool = True,
+        enable_prefill: bool = True,
+    ) -> None:
+        self.supports_packed_decode = enable_decode
+        self._packed_decode = None
+        self._replayssm_decode = None
+        self._chunk_kda = None
+        if enable_decode or enable_prefill:
+            try:
+                import helion  # noqa: F401
+            except ModuleNotFoundError as error:
+                if error.name != "helion":
+                    raise
+                raise ImportError(
+                    "The Helion package is required when a KDA backend is set to "
+                    "Helion. Install it with: pip install helion==1.4.0"
+                ) from None
+        if enable_decode:
+            from sglang.kernels.ops.attention.helion.kda_decode import (
+                helion_fused_recurrent_kda_packed_decode,
+            )
+            from sglang.kernels.ops.attention.helion.kda_replayssm import (
+                helion_fused_recurrent_kda_replayssm_decode,
+            )
+
+            self._packed_decode = helion_fused_recurrent_kda_packed_decode
+            self._replayssm_decode = helion_fused_recurrent_kda_replayssm_decode
+        if enable_prefill:
+            from sglang.kernels.ops.attention.helion.kda_prefill import chunk_kda
+
+            self._chunk_kda = chunk_kda
+        self._triton = triton_fallback or TritonKDAKernel()
+
+    def packed_decode(
+        self,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        *,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        scale: float,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        num_v_heads: int,
+        head_v_dim: int,
+        lower_bound: float | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        assert self._packed_decode is not None
+        batch_size = mixed_qkv.shape[0]
+        out = mixed_qkv.new_empty(batch_size, 1, num_v_heads, head_v_dim)
+        replayssm_d = kwargs.get("replayssm_d")
+        replayssm_k = kwargs.get("replayssm_k")
+        replayssm_g = kwargs.get("replayssm_g")
+        replayssm_write_pos = kwargs.get("replayssm_write_pos")
+        if (
+            replayssm_d is not None
+            and replayssm_k is not None
+            and replayssm_g is not None
+            and replayssm_write_pos is not None
+        ):
+            assert self._replayssm_decode is not None
+            self._replayssm_decode(
+                mixed_qkv=mixed_qkv,
+                a=a.reshape(batch_size, num_v_heads, -1).contiguous(),
+                b=b.reshape(batch_size, num_v_heads).contiguous(),
+                A_log=A_log.reshape(-1),
+                dt_bias=dt_bias.reshape(num_v_heads, -1).contiguous(),
+                scale=scale,
+                initial_state=ssm_states,
+                d_cache=replayssm_d,
+                k_cache=replayssm_k,
+                g_cache=replayssm_g,
+                out=out,
+                ssm_state_indices=cache_indices,
+                write_pos=replayssm_write_pos,
+                force_flush=kwargs.get("replayssm_force_flush"),
+                use_qk_l2norm_in_kernel=True,
+                lower_bound=lower_bound,
+            )
+            return out.transpose(0, 1)
+
+        if a.ndim != 2:
+            a = a.reshape(batch_size, -1)
+        if b.ndim != 2:
+            b = b.reshape(batch_size, -1)
+        self._packed_decode(
+            mixed_qkv=mixed_qkv,
+            a=a,
+            b=b,
+            A_log=A_log.reshape(-1),
+            dt_bias=dt_bias.reshape(-1),
+            scale=scale,
+            initial_state=ssm_states,
+            out=out,
+            ssm_state_indices=cache_indices,
+            use_qk_l2norm_in_kernel=True,
+            lower_bound=lower_bound,
+        )
+        return out.transpose(0, 1)
+
+    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 self._triton.decode(
+            q,
+            k,
+            v,
+            a,
+            b,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            **kwargs,
+        )
+
+    def extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        *,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        A_log: torch.Tensor | None = None,
+        dt_bias: torch.Tensor | None = None,
+        lower_bound: float | None = None,
+        return_intermediate_states: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        assert self._chunk_kda is not None
+        return self._chunk_kda(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            initial_state=ssm_states,
+            initial_state_indices=cache_indices,
+            use_qk_l2norm_in_kernel=True,
+            cu_seqlens=query_start_loc,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            lower_bound=lower_bound,
+            output_intermediate_states=return_intermediate_states,
+        )

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

@@ -20,6 +20,7 @@ class LinearAttnKernelBackend(Enum):
     FLASHKDA = "flashkda"
     NVIDIA_KDA = "nvidia_kda"
     PTX_KDA = "ptx_kda"
+    HELION = "helion"
     CUSTOM = "custom"
 
     @classmethod
@@ -47,6 +48,9 @@ class LinearAttnKernelBackend(Enum):
     def is_ptx_kda(self):
         return self == LinearAttnKernelBackend.PTX_KDA
 
+    def is_helion(self):
+        return self == LinearAttnKernelBackend.HELION
+
     def is_custom(self):
         return self == LinearAttnKernelBackend.CUSTOM
 

python/sglang/srt/server_args.py

@@ -377,6 +377,7 @@ LINEAR_ATTN_KERNEL_BACKEND_CHOICES = [
     "flashkda",
     "nvidia_kda",
     "ptx_kda",
+    "helion",
 ]
 
 
@@ -2571,7 +2572,7 @@ class ServerArgs:
     linear_attn_backend: A[
         str,
         Arg(
-            help="The default kernel backend for linear attention (GDN/KDA). Can be overridden per-mode by --linear-attn-decode-backend and --linear-attn-prefill-backend.",
+            help="The default kernel backend for linear attention (GDN/KDA). Can be overridden per-mode by --linear-attn-decode-backend and --linear-attn-prefill-backend. The Helion backend is KDA-only.",
             choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
         ),
         NS("exec.mamba"),
@@ -2606,11 +2607,10 @@ class ServerArgs:
         bool,
         "Enable the ReplaySSM buffered output-only linear-attn decode kernel. "
         "Primarily a GDN (scalar-gate) decode-bandwidth optimization (~1.2-1.5x "
-        "at batch >= 64). The unified kernel also supports KDA (per-K gate) and "
+        "at batch >= 64). KDA uses its selected Triton or Helion implementation, "
-        "is numerically correct, but KDA decode is SLOWER than the packed "
+        "but its per-K gate ring is larger and ReplaySSM is typically slower "
-        "baseline (the per-K g_cache is K x larger and the reconstruction "
+        "than packed KDA decode; benchmark before enabling it. Requires the "
-        "refolds the per-K decay every step), so it is not recommended for KDA "
+        "Triton linear-attn decode backend, or Helion for KDA, and "
-        "models. Requires the Triton linear-attn decode backend and "
         "--mamba-radix-cache-strategy no_buffer (the default).",
         NS("exec.mamba"),
     ] = False
@@ -6171,6 +6171,7 @@ class ServerArgs:
         # Fixed in FlashInfer v0.6.7: flashinfer-ai/flashinfer#2810
         if (
             self.linear_attn_decode_backend is None
+            and self.linear_attn_backend != "helion"
             and is_sm100_supported()
             and self.mamba_ssm_dtype == "bfloat16"
             # Stage 4: flashinfer's recurrent_kda compiles the state slot stride
@@ -6249,8 +6250,8 @@ class ServerArgs:
                 f"got CUDA {cuda_version or 'unknown'}"
             )
 
-        # GDN ReplaySSM buffered decode guards. Runs on the Triton GDN decode
+        # ReplaySSM buffered decode guards. Runs on Triton, or Helion for KDA.
-        # backend. cuda-graph is supported (slice 1b: CUDA-graph-safe static
+        # cuda-graph is supported (slice 1b: CUDA-graph-safe static
         # write-cursor buffers). The RADIX prefix cache is now supported (slice
         # 2b: the decode kernel force-flushes the ring into temporal[slot] on
         # the radix track boundary `seq_lens % mamba_track_interval == 0`, and
@@ -6264,10 +6265,10 @@ class ServerArgs:
         # cursor of the donated/kept slot would not be reset there. Handling
         # that donation path is a follow-up; for now require no_buffer.
         if self.enable_linear_replayssm:
-            if decode != "triton":
+            if decode not in {"triton", "helion"}:
                 raise ValueError(
-                    "--enable-linear-replayssm requires the Triton "
+                    "--enable-linear-replayssm requires Triton, or Helion for "
-                    "linear-attn decode backend, got "
+                    "KDA, as the linear-attn decode backend; got "
                     f"--linear-attn-decode-backend={decode!r}."
                 )
             from sglang.srt.arg_groups.overrides import (
@@ -8285,21 +8286,20 @@ class ServerArgs:
         # The Mamba/KDA state is stored in envelope-strided views; only
         # stride-audited kernels may read it (Stage 4 audit, per slot):
         # - decode: triton; flashinfer (recurrent_kda compiles the state slot
-        #   stride as a free int64 — natively strided); cutedsl (KDA fused
+        #   stride as a free int64); helion (specializes KDA state strides 0-3
-        #   sigmoid-gating update made stride-safe) on KDA-hybrid models only —
+        #   and rejects a non-unit innermost stride); cutedsl (KDA fused sigmoid-
-        #   cutedsl_gdn still compiles h0 against a contiguous dummy.
+        #   gating update is stride-safe) on KDA-hybrid models only.
-        # - prefill: triton; flashkda (wrapper gathers/scatters a contiguous
+        # - prefill: triton; flashkda (the wrapper gathers/scatters a contiguous
-        #   per-slot copy, external kernel never sees the pool); cutedsl
+        #   per-slot copy); helion; cutedsl (kernel_h compiles h0/ht with dynamic
-        #   (kernel_h compiles h0/ht with dynamic int64 strides), same
+        #   int64 strides), with the same KDA-only caveat.
-        #   KDA-only caveat.
         # - mamba (mamba2/short-conv state): triton only.
         # use_mla_backend() distinguishes the KDA-hybrid family (K3/KimiLinear
-        # are MLA-hybrid) from GDN models (GQA-hybrid) for the cutedsl caveat.
+        # are MLA-hybrid) from GDN models (GQA-hybrid) for the KDA-only caveat.
         decode_allowed = {"triton", "flashinfer"}
         prefill_allowed = {"triton", "flashkda"}
         if self.use_mla_backend():
-            decode_allowed.add("cutedsl")
+            decode_allowed.update({"cutedsl", "helion"})
-            prefill_allowed.add("cutedsl")
+            prefill_allowed.update({"cutedsl", "helion"})
         resolved_linear_decode = (
             self.linear_attn_decode_backend or self.linear_attn_backend
         )

test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py

@@ -217,6 +217,19 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
         tree_verify.assert_called_once()
         flashinfer_kernel.target_verify.assert_not_called()
 
+    def test_helion_backend_reports_kda_only(self):
+        cases = (
+            (LinearAttnKernelBackend.HELION, LinearAttnKernelBackend.TRITON),
+            (LinearAttnKernelBackend.TRITON, LinearAttnKernelBackend.HELION),
+        )
+        for decode_backend, prefill_backend in cases:
+            with self.subTest(
+                decode_backend=decode_backend,
+                prefill_backend=prefill_backend,
+            ):
+                with self.assertRaisesRegex(ValueError, "supports KDA only"):
+                    GDNKernelDispatcher(decode_backend, prefill_backend)
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/layers/attention/test_kda_helion_dispatcher.py

@@ -0,0 +1,197 @@
+import unittest
+from unittest.mock import ANY, MagicMock, patch
+
+import torch
+
+from sglang.srt.layers.attention.linear.kda_backend import KDAKernelDispatcher
+from sglang.srt.layers.attention.linear.kernels.kda_helion import HelionKDAKernel
+from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
+from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestHelionKDADispatcher(unittest.TestCase):
+    def _make_dispatcher(self, decode_backend, prefill_backend):
+        helion_kernel = MagicMock(supports_packed_decode=True)
+        with (
+            patch(
+                "sglang.srt.layers.attention.linear.kda_backend.is_cuda",
+                return_value=True,
+            ),
+            patch(
+                "sglang.srt.layers.attention.linear.kernels.kda_helion."
+                "HelionKDAKernel",
+                return_value=helion_kernel,
+            ) as constructor,
+        ):
+            dispatcher = KDAKernelDispatcher(
+                decode_backend=decode_backend,
+                prefill_backend=prefill_backend,
+                verify_backend=LinearAttnKernelBackend.TRITON,
+            )
+        return dispatcher, helion_kernel, constructor
+
+    def test_combined_backend_reuses_adapter_and_keeps_triton_verify(self):
+        dispatcher, helion_kernel, constructor = self._make_dispatcher(
+            LinearAttnKernelBackend.HELION,
+            LinearAttnKernelBackend.HELION,
+        )
+
+        constructor.assert_called_once_with(
+            triton_fallback=ANY,
+            enable_decode=True,
+            enable_prefill=True,
+        )
+        self.assertIs(dispatcher.decode_kernel, helion_kernel)
+        self.assertIs(dispatcher.extend_kernel, helion_kernel)
+        self.assertIsInstance(dispatcher.verify_kernel, TritonKDAKernel)
+        self.assertTrue(dispatcher.supports_packed_decode)
+
+    def test_decode_only_keeps_triton_prefill_and_verify(self):
+        dispatcher, helion_kernel, constructor = self._make_dispatcher(
+            LinearAttnKernelBackend.HELION,
+            LinearAttnKernelBackend.TRITON,
+        )
+
+        constructor.assert_called_once_with(
+            triton_fallback=ANY,
+            enable_decode=True,
+            enable_prefill=False,
+        )
+        self.assertIs(dispatcher.decode_kernel, helion_kernel)
+        self.assertIsInstance(dispatcher.extend_kernel, TritonKDAKernel)
+        self.assertIsInstance(dispatcher.verify_kernel, TritonKDAKernel)
+
+    def test_prefill_only_keeps_triton_decode_and_verify(self):
+        dispatcher, helion_kernel, constructor = self._make_dispatcher(
+            LinearAttnKernelBackend.TRITON,
+            LinearAttnKernelBackend.HELION,
+        )
+
+        constructor.assert_called_once_with(
+            triton_fallback=ANY,
+            enable_decode=False,
+            enable_prefill=True,
+        )
+        self.assertIsInstance(dispatcher.decode_kernel, TritonKDAKernel)
+        self.assertIs(dispatcher.extend_kernel, helion_kernel)
+        self.assertIsInstance(dispatcher.verify_kernel, TritonKDAKernel)
+
+    def test_enum_recognizes_helion(self):
+        backend = LinearAttnKernelBackend("helion")
+        self.assertIs(backend, LinearAttnKernelBackend.HELION)
+        self.assertTrue(backend.is_helion())
+
+    def test_replayssm_decode_uses_native_helion_kernel(self):
+        kernel = HelionKDAKernel.__new__(HelionKDAKernel)
+        kernel._packed_decode = MagicMock()
+        kernel._replayssm_decode = MagicMock()
+        kernel._triton = MagicMock()
+        mixed_qkv = torch.empty(2, 20)
+        a = torch.empty(2, 8)
+        b = torch.empty(2, 1)
+        state = torch.empty(2, 1, 4, 8)
+        indices = torch.arange(2, dtype=torch.int32)
+        force_flush = torch.zeros(2, dtype=torch.int32)
+        replay_args = {
+            "replayssm_d": torch.empty(2, 1, 4, 4),
+            "replayssm_k": torch.empty(2, 1, 4, 8),
+            "replayssm_g": torch.empty(2, 1, 4, 8),
+            "replayssm_write_pos": torch.zeros(2, dtype=torch.int32),
+            "replayssm_force_flush": force_flush,
+        }
+
+        result = kernel.packed_decode(
+            mixed_qkv,
+            a,
+            b,
+            A_log=torch.empty(1),
+            dt_bias=torch.empty(8),
+            scale=0.5,
+            ssm_states=state,
+            cache_indices=indices,
+            num_v_heads=1,
+            head_v_dim=4,
+            lower_bound=-5.0,
+            **replay_args,
+        )
+
+        self.assertEqual(result.shape, (1, 2, 1, 4))
+        kernel._packed_decode.assert_not_called()
+        kernel._replayssm_decode.assert_called_once()
+        self.assertIs(
+            kernel._replayssm_decode.call_args.kwargs["force_flush"], force_flush
+        )
+        self.assertEqual(kernel._replayssm_decode.call_args.kwargs["lower_bound"], -5.0)
+        kernel._triton.packed_decode.assert_not_called()
+
+    def test_packed_decode_forwards_lower_bound(self):
+        kernel = HelionKDAKernel.__new__(HelionKDAKernel)
+        kernel._packed_decode = MagicMock()
+        kernel._triton = MagicMock()
+        mixed_qkv = torch.empty(2, 16)
+        a = torch.empty(2, 8)
+        b = torch.empty(2, 1)
+        a_log = torch.empty(1)
+        dt_bias = torch.empty(8)
+        state = torch.empty(2, 1, 4, 8)
+        indices = torch.arange(2, dtype=torch.int32)
+
+        kernel.packed_decode(
+            mixed_qkv,
+            a,
+            b,
+            A_log=a_log,
+            dt_bias=dt_bias,
+            scale=0.5,
+            ssm_states=state,
+            cache_indices=indices,
+            num_v_heads=1,
+            head_v_dim=4,
+            lower_bound=-5.0,
+        )
+
+        kernel._packed_decode.assert_called_once()
+        self.assertEqual(kernel._packed_decode.call_args.kwargs["lower_bound"], -5.0)
+
+    def test_replayssm_accepts_helion_and_rejects_other_backends(self):
+        with (
+            patch("sglang.srt.server_args.is_sm100_supported", return_value=False),
+            patch("sglang.srt.server_args.is_cuda", return_value=False),
+        ):
+            helion_args = ServerArgs(
+                model_path="dummy",
+                linear_attn_decode_backend="helion",
+                enable_linear_replayssm=True,
+            )
+            helion_args._handle_linear_attn_backend()
+
+            flashinfer_args = ServerArgs(
+                model_path="dummy",
+                linear_attn_decode_backend="flashinfer",
+                enable_linear_replayssm=True,
+            )
+            with self.assertRaisesRegex(ValueError, "Triton, or Helion"):
+                flashinfer_args._handle_linear_attn_backend()
+
+    def test_explicit_base_backend_is_not_replaced_by_flashinfer(self):
+        args = ServerArgs(
+            model_path="dummy",
+            linear_attn_backend="helion",
+            mamba_ssm_dtype="bfloat16",
+        )
+        with (
+            patch("sglang.srt.server_args.is_sm100_supported", return_value=True),
+            patch("sglang.srt.server_args.is_cuda", return_value=False),
+        ):
+            args._handle_linear_attn_backend()
+
+        self.assertIsNone(args.linear_attn_decode_backend)
+        self.assertEqual(args.linear_attn_backend, "helion")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/server_args/test_page_major_backend_allowlist.py

@@ -37,7 +37,14 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
-def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
+def _accepts(
+    backend: str,
+    *,
+    use_mla: bool,
+    unified: bool = True,
+    linear_decode: str | None = None,
+    linear_prefill: str | None = None,
+) -> bool:
     """Run just `_handle_page_major_kv_layout` against a minimal stand-in.
 
     ServerArgs' real constructor pulls in a model config; this exercises the
@@ -53,8 +60,8 @@ def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
         "prefill_attention_backend": None,
         "decode_attention_backend": None,
         "linear_attn_backend": "triton",
-        "linear_attn_decode_backend": None,
+        "linear_attn_decode_backend": linear_decode,
-        "linear_attn_prefill_backend": None,
+        "linear_attn_prefill_backend": linear_prefill,
         "mamba_backend": "triton",
     }.items():
         object.__setattr__(sa, name, value)
@@ -115,6 +122,17 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
                     f"{backend} has no dense-id remapping and must be rejected",
                 )
 
+    def test_helion_linear_attention_is_kda_only(self):
+        for unified in (True, False):
+            for phase in ("decode", "prefill"):
+                kwargs = {f"linear_{phase}": "helion"}
+                self.assertTrue(
+                    _accepts("triton", use_mla=True, unified=unified, **kwargs)
+                )
+                self.assertFalse(
+                    _accepts("triton", use_mla=False, unified=unified, **kwargs)
+                )
+
 
 if __name__ == "__main__":
     unittest.main()