[GDN][KDA] ReplaySSM buffered output-only decode for Linear Attention (#28451)

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

python/sglang/srt/configs/mamba_utils.py

@@ -124,6 +124,13 @@ class BaseLinearStateParams(ABC):
             + ssm_numel * self.dtype.temporal.itemsize
         ) * len(self.layers)
 
+    @property
+    def is_kda(self) -> bool:
+        """KDA per-K-channel gate vs GDN/Mamba2 per-head scalar gate. Selects
+        the ReplaySSM ring ``g_cache`` layout ([.., L] scalar vs [.., L, K]
+        per-K) and the gate-generic decode kernel's ``IS_KDA`` path."""
+        return False
+
 
 @dataclass(kw_only=True, frozen=True)
 class Mamba2StateShape:
@@ -137,6 +144,10 @@ class Mamba2StateShape:
     head_dim: int
     state_size: int
     conv_kernel: int
+    # Number of key/group heads after TP sharding (== runtime `H` the packed
+    # GDN kernels infer from `mixed_qkv`). Used by the GDN ReplaySSM ring
+    # buffer (k_cache) to size/stride exactly like the kernel expects.
+    num_k_heads_per_tp: int = 1
 
     @staticmethod
     def create(
@@ -149,6 +160,16 @@ class Mamba2StateShape:
         state_size: int,
         conv_kernel: int,
     ) -> "Mamba2StateShape":
+        # The q/k projections are sharded by `num_k_heads // tp` heads (the
+        # ORIGINAL n_groups, before the conv head-shard extension below), so the
+        # runtime `H` the packed kernels see equals divide(n_groups, tp). Only
+        # meaningful (and only consumed) for the GDN ReplaySSM path, which
+        # requires evenly divisible heads; fall back to ceil-div otherwise.
+        num_k_heads_per_tp = (
+            divide(n_groups, tp_world_size)
+            if n_groups % tp_world_size == 0
+            else -(-n_groups // tp_world_size)
+        )
         # if n_groups is not divisible by world_size, need to extend the shards
         # to ensure all groups needed by a head is sharded along with it
         if n_groups % tp_world_size != 0:
@@ -174,6 +195,7 @@ class Mamba2StateShape:
             head_dim=head_dim,
             state_size=state_size,
             conv_kernel=conv_kernel,
+            num_k_heads_per_tp=num_k_heads_per_tp,
         )
 
 
@@ -193,6 +215,10 @@ class KimiLinearStateShape:
     head_k_dim: int
     conv_kernel: int
     num_spec: int
+    # Number of key heads after TP sharding (== runtime ``H`` the KDA packed
+    # kernels infer from ``mixed_qkv``). Mirrors Mamba2StateShape; consumed by
+    # the ReplaySSM ring (k_cache) to size/stride exactly like the kernel.
+    num_k_heads_per_tp: int = 1
 
     @staticmethod
     def create(
@@ -209,6 +235,11 @@ class KimiLinearStateShape:
             num_k_heads = num_heads
         if head_k_dim is None:
             head_k_dim = head_dim
+        num_k_heads_per_tp = (
+            divide(num_k_heads, tp_world_size)
+            if num_k_heads % tp_world_size == 0
+            else -(-num_k_heads // tp_world_size)
+        )
 
         proj_size = num_heads * head_dim
         proj_k_size = num_k_heads * head_k_dim
@@ -231,9 +262,14 @@ class KimiLinearStateShape:
             head_k_dim=head_k_dim,
             conv_kernel=conv_kernel_size,
             num_spec=num_spec,
+            num_k_heads_per_tp=num_k_heads_per_tp,
         )
 
 
 @dataclass(kw_only=True, frozen=True)
 class KimiLinearCacheParams(BaseLinearStateParams):
     shape: KimiLinearStateShape
+
+    @property
+    def is_kda(self) -> bool:
+        return True

python/sglang/srt/layers/attention/fla/bench_gdn_replayssm_decode.py

@@ -0,0 +1,166 @@
+"""Microbenchmark: buffered output-only GDN decode (ReplaySSM Part A) vs. the
+existing packed GDN decode kernel.
+
+Compares per-step decode latency of
+``fused_recurrent_gated_delta_rule_packed_decode`` (writes the full recurrent
+state S every step) against ``fused_recurrent_gdn_replayssm_decode`` at
+L in {1, 8, 16} (writes the full state only every L steps) across batch sizes
+{1, 16, 64, 256} for a realistic GDN config (HV=32, K=V=128).
+
+The win is per-step HBM *state* traffic: the packed kernel reads + writes S
+(~2 * num_slots * HV * V * K * 4 bytes / step for an fp32 state), while the
+ReplaySSM kernel reads S every step but writes it only 1-in-L steps, plus a
+small ring append (d:[HV,V], k:[H,K], g:[HV] per step). The amortized state
+traffic ratio is reported per L.
+
+Run::
+
+    python -m sglang.srt.layers.attention.fla.bench_gdn_replayssm_decode
+
+Requires a GPU (Triton).
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+import triton
+
+from sglang.srt.layers.attention.fla.fused_recurrent import (
+    fused_recurrent_gated_delta_rule_packed_decode,
+)
+from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+    fused_recurrent_gdn_replayssm_decode,
+)
+
+
+def _make_static(B, H, HV, K, V, dtype, device):
+    qk_dim = 2 * H * K
+    v_dim = HV * V
+    mixed_qkv = torch.randn(B, qk_dim + v_dim, device=device, dtype=dtype)
+    a = torch.randn(B, HV, device=device, dtype=dtype) * 0.5
+    b = torch.randn(B, HV, device=device, dtype=dtype)
+    A_log = (torch.randn(HV, device=device, dtype=torch.float32) * 0.3).contiguous()
+    dt_bias = (torch.randn(HV, device=device, dtype=torch.float32) * 0.1).contiguous()
+    return mixed_qkv, a, b, A_log, dt_bias
+
+
+def _state_bytes_per_step(B, HV, K, V, L, dtype):
+    """Amortized per-step HBM *state* traffic (bytes), state in fp32.
+
+    packed: read S + write S every step.
+    replay: read S every step; write S once per L steps; append ring records
+            (d:[HV,V] in `dtype`, k:[H,K] in `dtype` shared across HV//H, g:[HV]
+            fp32) every step. We report the dominant fp32-state terms; ring
+            appends are tiny by comparison and shown separately.
+    """
+    fp32 = 4
+    state_elems = B * HV * V * K  # one record per active request slot
+    packed = (state_elems * fp32) * 2  # read + write
+    replay = (state_elems * fp32) * (1 + 1.0 / L)  # read every step + write 1/L
+    return packed, replay
+
+
+def _bench_cfg(B, H, HV, K, V, Ls, dtype, device, num_slots=None, warmup=25, rep=100):
+    num_slots = num_slots or B
+    mixed_qkv, a, b, A_log, dt_bias = _make_static(B, H, HV, K, V, dtype, device)
+    scale = K**-0.5
+    cache_indices = torch.arange(B, device=device, dtype=torch.int32)
+
+    # packed decode
+    state = torch.randn(num_slots, HV, V, K, device=device, dtype=torch.float32)
+    out = mixed_qkv.new_empty(B, 1, HV, V)
+
+    def run_packed():
+        fused_recurrent_gated_delta_rule_packed_decode(
+            mixed_qkv=mixed_qkv,
+            a=a,
+            b=b,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            scale=scale,
+            initial_state=state,
+            out=out,
+            ssm_state_indices=cache_indices,
+            use_qk_l2norm_in_kernel=True,
+        )
+
+    t_packed = triton.testing.do_bench(run_packed, warmup=warmup, rep=rep)
+
+    rows = []
+    for L in Ls:
+        rstate = torch.randn(num_slots, HV, V, K, device=device, dtype=torch.float32)
+        d_cache = torch.zeros(num_slots, HV, L, V, device=device, dtype=dtype)
+        k_cache = torch.zeros(num_slots, H, L, K, device=device, dtype=dtype)
+        g_cache = torch.zeros(num_slots, HV, L, device=device, dtype=torch.float32)
+        write_pos = torch.zeros(B, device=device, dtype=torch.int32)
+        rout = mixed_qkv.new_empty(B, 1, HV, V)
+        nk = 1 if L == 1 else 2
+
+        def run_replay():
+            fused_recurrent_gdn_replayssm_decode(
+                mixed_qkv=mixed_qkv,
+                a=a,
+                b=b,
+                A_log=A_log,
+                dt_bias=dt_bias,
+                scale=scale,
+                initial_state=rstate,
+                d_cache=d_cache,
+                k_cache=k_cache,
+                g_cache=g_cache,
+                out=rout,
+                ssm_state_indices=cache_indices,
+                write_pos=write_pos,
+                use_qk_l2norm_in_kernel=True,
+                nk=nk,
+            )
+
+        t_replay = triton.testing.do_bench(run_replay, warmup=warmup, rep=rep)
+        packed_bytes, replay_bytes = _state_bytes_per_step(B, HV, K, V, L, dtype)
+        rows.append((L, t_replay, t_packed / t_replay, replay_bytes / packed_bytes))
+    return t_packed, rows
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--hv", type=int, default=32, help="num value heads")
+    parser.add_argument("--h", type=int, default=16, help="num key/query heads")
+    parser.add_argument("--k", type=int, default=128)
+    parser.add_argument("--v", type=int, default=128)
+    parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 16, 64, 256])
+    parser.add_argument("--ls", type=int, nargs="+", default=[1, 8, 16])
+    parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16")
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA / Triton required for this microbenchmark.")
+    device = "cuda"
+    dtype = {
+        "bf16": torch.bfloat16,
+        "fp16": torch.float16,
+        "fp32": torch.float32,
+    }[args.dtype]
+
+    print(
+        f"GDN ReplaySSM decode microbench  HV={args.hv} H={args.h} "
+        f"K={args.k} V={args.v} dtype={args.dtype}\n"
+        "per-step latency (ms); speedup = packed/replay; "
+        "state-traffic = replay/packed (lower is better)"
+    )
+    for B in args.batch_sizes:
+        t_packed, rows = _bench_cfg(
+            B, args.h, args.hv, args.k, args.v, args.ls, dtype, device
+        )
+        print(f"\nB={B:<4d}  packed={t_packed:.4f} ms")
+        for L, t_replay, speedup, traffic_ratio in rows:
+            print(
+                f"    L={L:<3d} replay={t_replay:.4f} ms  "
+                f"speedup={speedup:5.2f}x  "
+                f"state-traffic={traffic_ratio:5.2f}x"
+            )
+
+
+if __name__ == "__main__":
+    main()

python/sglang/srt/layers/attention/fla/fused_recurrent_linear_replayssm.py

@@ -0,0 +1,628 @@
+# Buffered output-only linear-attention decode (ReplaySSM Part A), ported to
+# SGLang. Covers BOTH gate granularities with one kernel:
+#   * GDN  (``IS_KDA=False``): per-head SCALAR gate ``alpha = exp(g)``.
+#   * KDA  (``IS_KDA=True``):  per-K-channel gate ``alpha[k] = exp(g[k])`` —
+#     the state decays column-wise, ``S' = S . Diag(alpha) + d k^T``.
+# GDN is the special case of KDA with all per-K decays equal; a single
+# ``IS_KDA`` constexpr selects the gate path and the GDN path is bit-for-bit
+# the original (no regression to the validated GDN kernel).
+#
+# This is a STANDALONE increment: kernel + wrapper only. It is NOT yet wired
+# into the memory pool / radix cache / scheduler / backend dispatch. The caller
+# (currently the correctness test and the microbenchmark) owns the ring tensors.
+#
+# Idea (vs. ``fused_recurrent_gated_delta_rule_packed_decode`` / ``..._kda_...``):
+#   The plain packed decode reads the full recurrent state S [HV, V, K] from
+#   HBM and writes it back *every* decode step (~8*d*n bytes/step of state
+#   traffic for an fp32 state, read+write).  ReplaySSM keeps a small per-slot
+#   ring buffer of the last L steps' (d, k, g) and only WRITES the full state
+#   every L steps (a "flush"); on non-flush steps it appends a tiny
+#   (d, k, g) record and reconstructs the readout from the checkpoint S0 plus
+#   the buffer.  S0 is still READ every step, so per-step state traffic drops
+#   from read+write (~8*d*n) to read-only (~4*d*n) -> roughly halved.
+#
+# Math (single head, single step; matches the packed decode kernels exactly).
+# Let ``a = exp(g)`` be the decay (scalar for GDN, per-K vector for KDA) and
+# ``S`` the state *before* this token:
+#   d_cur = beta * (v - (S . Diag(a)) . k) = beta * (v - S . (a (.) k))
+#   o     = (S . Diag(a)) . q + d_cur*(k^T q) = S . (a (.) q) + d_cur*(k^T q)
+#   S_new = S . Diag(a) + d_cur k^T          # only persisted on flush
+# where ``(.)`` is elementwise over K.  For GDN ``a`` is scalar so
+# ``S . (a (.) q) = a * (S . q)`` (the cheap scalar post-multiply); for KDA the
+# per-K ``a`` folds into q/k before the matvec.  k^T q uses the RAW current
+# k/q (the rank-1 term), so it is identical for both gate types.
+#
+# Buffered reconstruction: with buffered steps j=0..m-1 holding (d_j, k_j, g_j),
+# the state *before* the current token is
+#   S = Diag(A) . S0 + sum_j d_j (W_j (.) k_j)^T          (per-K form)
+# with A[c]   = exp(sum_j g_j[c])                  (total decay, per-K)
+#      W_j[c] = exp(sum_i g_i[c] - cumsum_inclusive_j[c]) = prod_{i>j} a_i[c].
+# For GDN A and W_j are scalars (g is K-independent) and W_j folds onto d_j
+# instead of k_j (either factor works for a scalar). S is reconstructed in
+# K-tiles and immediately read with q (and k) -> the [V,K] state tile is never
+# fully materialized to HBM on a non-flush step.
+#
+# At L=1 the ring is always empty and ``write_pos == L-1`` every step, so the
+# reconstruction term is zero, the total decay is 1, and this kernel reduces
+# *algebraically* to the corresponding packed-decode kernel.
+#
+# SPDX-License-Identifier: Apache-2.0
+# Ported from vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py
+# (ReplaySSM, commit 3c85112) and adapted to SGLang's packed GDN/KDA decode
+# layout.
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def fused_recurrent_linear_replayssm_decode_kernel(
+    mixed_qkv,  # [B, 2*H*K + HV*V] packed (q | k | v) after conv1d
+    a,  # GDN: [B, HV] gate input ; KDA: [B, HV, K] per-K gate input
+    b,  # [B, HV] beta input b
+    A_log,  # [HV] log-space decay parameter (per-head scalar, both gate types)
+    dt_bias,  # GDN: [HV] ; KDA: [HV, K] time-step bias
+    o,  # [B, HV, V] output (written every step)
+    h0,  # [num_slots, HV, V, K] checkpoint state (read every step)
+    ht,  # [num_slots, HV, V, K] checkpoint state (written only on flush; == h0)
+    d_cache,  # [num_slots, HV, L, V] ring: corrected delta vectors
+    k_cache,  # [num_slots, H,  L, K] ring: (normed/scaled) keys
+    g_cache,  # GDN: [num_slots, HV, L] ; KDA: [num_slots, HV, L, K] log-decay gates (fp32)
+    ssm_state_indices,  # [B] physical state slot per decode row
+    write_pos,  # [B] int32 per-row ring cursor (0..L-1)
+    force_flush,  # [B] int32: !=0 forces a flush this step (radix track boundary)
+    scale,
+    stride_mixed_qkv_tok: tl.constexpr,
+    stride_a_tok: tl.constexpr,
+    stride_b_tok: tl.constexpr,
+    stride_init_state_token: tl.constexpr,
+    stride_final_state_token: tl.constexpr,
+    stride_indices_seq: tl.constexpr,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    BC: tl.constexpr,
+    NK: tl.constexpr,
+    BKT: tl.constexpr,
+    MAX_CACHE_LEN: tl.constexpr,
+    SOFTPLUS_THRESHOLD: tl.constexpr,
+    USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+    HAS_FORCE_FLUSH: tl.constexpr,
+    IS_KDA: tl.constexpr,
+):
+    i_v = tl.program_id(0)
+    i_n = tl.program_id(1)
+    i_hv = tl.program_id(2)
+    i_h = i_hv // (HV // H)
+
+    o_v = i_v * BV + tl.arange(0, BV)
+    o_c = tl.arange(0, BC)
+    mask_v = o_v < V
+
+    # Resolve the physical state slot; zero the output and bail for padded rows.
+    state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64)
+    p_o = o + (i_n * HV + i_hv) * V + o_v
+    if state_idx < 0:
+        tl.store(
+            p_o,
+            tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty),
+            mask=mask_v,
+        )
+        return
+
+    # Per-row buffer cursor and flush flag (device-side branch, no host branch),
+    # plus the set of valid (already committed) cache positions.
+    b_write_pos = tl.load(write_pos + i_n).to(tl.int64)
+    b_is_flush = b_write_pos == MAX_CACHE_LEN - 1
+    if HAS_FORCE_FLUSH:
+        # A radix track-boundary (or any caller-forced) flush folds the partial
+        # ring (the real `write_pos` entries, NOT L-1) + current token into the
+        # checkpoint so an external snapshot reads an up-to-date state. cache_valid
+        # below still uses the true write_pos, so only committed entries are read.
+        b_is_flush = b_is_flush | (tl.load(force_flush + i_n) != 0)
+    cache_valid = o_c < b_write_pos
+
+    # Gate for the current token.  beta is a per-head scalar for both gate
+    # types; A_log is a per-head scalar for both.  The decay g/alpha is a
+    # per-head scalar for GDN (computed here) and a per-K vector for KDA
+    # (computed per K-tile inside the loop, since it is K-indexed).
+    #   g = -exp(A_log) * softplus(a + dt_bias);  alpha = exp(g);  beta = sigmoid(b)
+    A_log_val = tl.load(A_log + i_hv).to(tl.float32)
+    b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32)
+    beta_val = tl.sigmoid(b_val).to(b.dtype.element_ty).to(tl.float32)
+    if not IS_KDA:
+        a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32)
+        dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32)
+        x = a_val + dt_bias_val
+        softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
+        g_val = -tl.exp(A_log_val) * softplus_x
+        alpha_val = tl.exp(g_val)
+
+        # Replay decay over the committed cache, from the cached per-step gates.
+        # b_replay_decay[j] = exp(sum_i g_i - cumsum_inclusive_j) = prod_{i>j} alpha_i
+        p_g_main = g_cache + (state_idx * HV + i_hv) * MAX_CACHE_LEN + o_c
+        b_g_all = tl.load(p_g_main, mask=cache_valid, other=0.0).to(tl.float32)
+        b_g_prefix = tl.cumsum(b_g_all, axis=0)
+        b_g_total = tl.sum(b_g_all, axis=0)
+        b_replay_decay = tl.where(cache_valid, tl.exp(b_g_total - b_g_prefix), 0.0)
+        b_total_decay = tl.exp(b_g_total)
+
+    # Cached corrected-delta vectors d (K-independent).  Layout
+    # d_cache[slot, hv, L, V] -> index [V, BC] tile.
+    p_d_main = d_cache + (
+        ((state_idx * HV + i_hv) * MAX_CACHE_LEN + o_c[None, :]) * V + o_v[:, None]
+    )
+    b_d_all = tl.load(
+        p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0
+    ).to(tl.float32)
+    # Cast the (d, k) reconstruction-dot operands to the I/O dtype so tl.dot
+    # runs on TENSOR CORES (bf16 tensor cores for bf16, TF32 for fp32). This is
+    # the ReplaySSM reference path and is performance-critical: the L-deep
+    # reconstruction is L x the baseline's rank-1 compute, which only stays
+    # hidden under the (reduced) memory traffic if it runs on tensor cores. An
+    # IEEE fp32 dot here disables tensor cores and makes the kernel ~10-20x
+    # slower. Tensor-core precision (~4e-4 TF32 / ~1e-3 bf16) is benign
+    # end-to-end (ReplaySSM bf16 GSM8K parity); the unit test uses
+    # tensor-core-realistic tolerances. At L=1 the buffer is empty so this dot
+    # is identically zero and the path stays bit-exact regardless of precision.
+    # GDN folds the (scalar) replay decay onto d here; KDA folds the (per-K)
+    # replay decay onto the cached keys inside the K-tile loop instead.
+    if not IS_KDA:
+        b_d_tc = (b_d_all * b_replay_decay[None, :]).to(
+            p_o.dtype.element_ty
+        )  # [BV, BC]
+    else:
+        b_d_tc = b_d_all.to(p_o.dtype.element_ty)  # [BV, BC]
+
+    # Current token value (for the delta-rule update).
+    v_off = (2 * H * K) + i_hv * V + o_v
+    b_v = tl.load(
+        mixed_qkv + i_n * stride_mixed_qkv_tok + v_off, mask=mask_v, other=0
+    ).to(tl.float32)
+
+    # Optional q/k L2 norm: full-vector reciprocal norms (computed, not kept).
+    if USE_QK_L2NORM_IN_KERNEL:
+        o_kf = tl.arange(0, BK)
+        mask_kf = o_kf < K
+        p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok
+        qf = tl.load(p_mix + i_h * K + o_kf, mask=mask_kf, other=0).to(tl.float32)
+        kf = tl.load(p_mix + H * K + i_h * K + o_kf, mask=mask_kf, other=0).to(
+            tl.float32
+        )
+        q_rnorm = 1.0 / tl.sqrt(tl.sum(qf * qf) + 1e-6)
+        k_rnorm = 1.0 / tl.sqrt(tl.sum(kf * kf) + 1e-6)
+    else:
+        q_rnorm = 1.0
+        k_rnorm = 1.0
+
+    # Reconstruct S from the checkpoint + cached (d, k) in K-tiles and read it
+    # with the current (scaled) q and k.  K-tiling keeps the per-program tile
+    # small so the full [V, K] state is never materialized.  Also append the
+    # current key chunk to the ring cache (non-flush only).
+    b_state_q = tl.zeros([BV], dtype=tl.float32)
+    b_state_k = tl.zeros([BV], dtype=tl.float32)
+    cur_kq = tl.zeros([1], dtype=tl.float32)
+    write_k = (not b_is_flush) and (i_v == 0) and (i_hv == i_h * (HV // H))
+    write_g_kda = IS_KDA and (not b_is_flush) and (i_v == 0)
+    for kk in range(NK):
+        o_kt = kk * BKT + tl.arange(0, BKT)
+        mask_kt = o_kt < K
+        p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok
+        q_c = (
+            tl.load(p_mix + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32)
+            * q_rnorm
+        )
+        k_c = (
+            tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to(
+                tl.float32
+            )
+            * k_rnorm
+        )
+        q_cs = q_c * scale
+        # Rank-1 output term uses the RAW current k/q (gate-independent).
+        cur_kq += tl.sum(k_c * q_cs)
+
+        # This K-tile of the state: S_tile = Diag(A_tile) S0_tile + d (.) (W (.) k_cache).
+        p_h0_c = (
+            h0
+            + state_idx * stride_init_state_token
+            + i_hv * V * K
+            + o_v[:, None] * K
+            + o_kt[None, :]
+        )
+        b_h0_c = tl.load(p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0).to(
+            tl.float32
+        )
+        p_k_c = (
+            k_cache
+            + ((state_idx * H + i_h) * MAX_CACHE_LEN + o_c[:, None]) * K
+            + o_kt[None, :]
+        )
+
+        if not IS_KDA:
+            b_k_all_c = tl.load(
+                p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0
+            ).to(p_o.dtype.element_ty)
+            b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_tc, b_k_all_c).to(tl.float32)
+            # GDN: scalar current-token decay applied after the loop.
+            q_eff = q_cs
+            k_eff = k_c
+        else:
+            # KDA per-K decay: load this tile's cached gates [BC, BKT], form the
+            # per-K total / replay decay, fold the replay decay onto the cached
+            # keys and the total decay onto S0.
+            p_g_c = (
+                g_cache
+                + ((state_idx * HV + i_hv) * MAX_CACHE_LEN + o_c[:, None]) * K
+                + o_kt[None, :]
+            )
+            b_g_all_c = tl.load(
+                p_g_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0.0
+            ).to(tl.float32)
+            b_g_prefix_c = tl.cumsum(b_g_all_c, axis=0)  # [BC, BKT]
+            b_g_total_c = tl.sum(b_g_all_c, axis=0)  # [BKT]
+            b_replay_decay_c = tl.where(
+                cache_valid[:, None],
+                tl.exp(b_g_total_c[None, :] - b_g_prefix_c),
+                0.0,
+            )  # [BC, BKT]
+            b_total_decay_c = tl.exp(b_g_total_c)  # [BKT]
+            b_k_all_c = tl.load(
+                p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0.0
+            ).to(tl.float32)
+            b_k_scaled = (b_k_all_c * b_replay_decay_c).to(p_o.dtype.element_ty)
+            b_h_c = b_h0_c * b_total_decay_c[None, :] + tl.dot(b_d_tc, b_k_scaled).to(
+                tl.float32
+            )
+            # KDA: current-token per-K decay folds into q/k for the readout.
+            p_a_c = a + i_n * stride_a_tok + i_hv * K + o_kt
+            p_dt_c = dt_bias + i_hv * K + o_kt
+            b_a_c = tl.load(p_a_c, mask=mask_kt, other=0.0).to(tl.float32)
+            b_dt_c = tl.load(p_dt_c, mask=mask_kt, other=0.0).to(tl.float32)
+            x_c = b_a_c + b_dt_c
+            softplus_c = tl.where(
+                x_c <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x_c)), x_c
+            )
+            g_cur_c = -tl.exp(A_log_val) * softplus_c  # [BKT]
+            alpha_cur_c = tl.exp(g_cur_c)
+            q_eff = q_cs * alpha_cur_c
+            k_eff = k_c * alpha_cur_c
+
+            # Append this tile of the current gate to the ring (non-flush only).
+            if write_g_kda:
+                p_cur_g = (
+                    g_cache
+                    + ((state_idx * HV + i_hv) * MAX_CACHE_LEN + b_write_pos) * K
+                    + o_kt
+                )
+                tl.store(p_cur_g, g_cur_c, mask=mask_kt & (b_write_pos < MAX_CACHE_LEN))
+
+        # Read the state with the (gate-folded) q and k, accumulated across tiles.
+        b_state_q += tl.sum(b_h_c * q_eff[None, :], axis=1)
+        b_state_k += tl.sum(b_h_c * k_eff[None, :], axis=1)
+
+        if write_k:
+            p_cur_k = (
+                k_cache
+                + ((state_idx * H + i_h) * MAX_CACHE_LEN + b_write_pos) * K
+                + o_kt
+            )
+            tl.store(
+                p_cur_k,
+                k_c.to(p_o.dtype.element_ty),
+                mask=mask_kt & (b_write_pos < MAX_CACHE_LEN),
+            )
+
+    # Current-token output: (S . Diag(a)) q + d_cur*(k . q), with the new
+    # corrected delta-rule vector d_cur = beta * (v - (S . Diag(a)) k).
+    # For GDN the per-head scalar decay is applied here; for KDA it was already
+    # folded into q/k above.
+    if not IS_KDA:
+        b_state_q *= alpha_val
+        b_state_k *= alpha_val
+    b_d_cur = beta_val * (b_v - b_state_k)
+    b_o = b_state_q + b_d_cur * tl.sum(cur_kq)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
+
+    if b_is_flush:
+        # Flush: fold the current token into the checkpoint, S_new = S.Diag(a) +
+        # d_cur k^T, and persist it.  Re-walk K chunks to rebuild S, then apply
+        # the update.  After this the ring is logically cleared (the caller
+        # resets write_pos to 0 on the next step).
+        for kk in range(NK):
+            o_kt = kk * BKT + tl.arange(0, BKT)
+            mask_kt = o_kt < K
+            p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok
+            k_c = (
+                tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to(
+                    tl.float32
+                )
+                * k_rnorm
+            )
+            p_h0_c = (
+                h0
+                + state_idx * stride_init_state_token
+                + i_hv * V * K
+                + o_v[:, None] * K
+                + o_kt[None, :]
+            )
+            b_h0_c = tl.load(
+                p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0
+            ).to(tl.float32)
+            p_k_c = (
+                k_cache
+                + ((state_idx * H + i_h) * MAX_CACHE_LEN + o_c[:, None]) * K
+                + o_kt[None, :]
+            )
+            if not IS_KDA:
+                b_k_all_c = tl.load(
+                    p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0
+                ).to(p_o.dtype.element_ty)
+                b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_tc, b_k_all_c).to(
+                    tl.float32
+                )
+                b_h_new_c = alpha_val * b_h_c + b_d_cur[:, None] * k_c[None, :]
+            else:
+                p_g_c = (
+                    g_cache
+                    + ((state_idx * HV + i_hv) * MAX_CACHE_LEN + o_c[:, None]) * K
+                    + o_kt[None, :]
+                )
+                b_g_all_c = tl.load(
+                    p_g_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0.0
+                ).to(tl.float32)
+                b_g_prefix_c = tl.cumsum(b_g_all_c, axis=0)
+                b_g_total_c = tl.sum(b_g_all_c, axis=0)
+                b_replay_decay_c = tl.where(
+                    cache_valid[:, None],
+                    tl.exp(b_g_total_c[None, :] - b_g_prefix_c),
+                    0.0,
+                )
+                b_total_decay_c = tl.exp(b_g_total_c)
+                b_k_all_c = tl.load(
+                    p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0.0
+                ).to(tl.float32)
+                b_k_scaled = (b_k_all_c * b_replay_decay_c).to(p_o.dtype.element_ty)
+                b_h_c = b_h0_c * b_total_decay_c[None, :] + tl.dot(
+                    b_d_tc, b_k_scaled
+                ).to(tl.float32)
+                p_a_c = a + i_n * stride_a_tok + i_hv * K + o_kt
+                p_dt_c = dt_bias + i_hv * K + o_kt
+                b_a_c = tl.load(p_a_c, mask=mask_kt, other=0.0).to(tl.float32)
+                b_dt_c = tl.load(p_dt_c, mask=mask_kt, other=0.0).to(tl.float32)
+                x_c = b_a_c + b_dt_c
+                softplus_c = tl.where(
+                    x_c <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x_c)), x_c
+                )
+                alpha_cur_c = tl.exp(-tl.exp(A_log_val) * softplus_c)  # [BKT]
+                b_h_new_c = (
+                    b_h_c * alpha_cur_c[None, :] + b_d_cur[:, None] * k_c[None, :]
+                )
+            p_ht_c = (
+                ht
+                + state_idx * stride_final_state_token
+                + i_hv * V * K
+                + o_v[:, None] * K
+                + o_kt[None, :]
+            )
+            tl.store(
+                p_ht_c,
+                b_h_new_c.to(p_ht_c.dtype.element_ty),
+                mask=mask_v[:, None] & mask_kt[None, :],
+            )
+    else:
+        # Non-flush: append the current token's corrected delta d to the cache
+        # (k chunks were written inside the loop; KDA's g chunks too).  GDN's
+        # scalar g is appended here.
+        p_cur_d = (
+            d_cache + ((state_idx * HV + i_hv) * MAX_CACHE_LEN + b_write_pos) * V + o_v
+        )
+        tl.store(
+            p_cur_d,
+            b_d_cur.to(p_cur_d.dtype.element_ty),
+            mask=mask_v & (b_write_pos < MAX_CACHE_LEN),
+        )
+        if (not IS_KDA) and (i_v == 0):
+            p_cur_g = g_cache + (state_idx * HV + i_hv) * MAX_CACHE_LEN + b_write_pos
+            tl.store(p_cur_g, g_val, mask=b_write_pos < MAX_CACHE_LEN)
+
+
+def fused_recurrent_linear_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,
+    is_kda: bool = False,
+    block_v: int | None = None,
+    num_warps: int = 1,
+    num_stages: int = 3,
+    nk: int = 2,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Buffered output-only linear-attention autoregressive decode (1 token/seq).
+
+    One kernel for both gate granularities, selected by ``is_kda``:
+      * ``is_kda=False`` (GDN): per-head SCALAR gate.  ``a``=[B, HV],
+        ``dt_bias``=[HV], ``g_cache``=[num_slots, HV, L].
+      * ``is_kda=True`` (KDA): per-K-channel gate.  ``a``=[B, HV, K],
+        ``dt_bias``=[HV, K], ``g_cache``=[num_slots, HV, L, K].
+    ``A_log`` is [HV] (per-head scalar) for both.
+
+    Same call surface as the packed decode plus the three ring caches
+    (``d_cache`` / ``k_cache`` / ``g_cache``) and the per-decode-row
+    ``write_pos`` cursor.  ``initial_state`` is both the checkpoint read (h0)
+    and the (flush-only) checkpoint write (ht), in place.
+
+    Allocates nothing persistent: the caller owns the ring tensors and is
+    responsible for advancing / resetting ``write_pos`` (e.g. ``(write_pos+1) %
+    L`` after each step).  This is a STANDALONE kernel; the memory-pool / cache
+    integration is a later phase.
+    """
+    if mixed_qkv.ndim != 2:
+        raise ValueError(f"`mixed_qkv` must be 2D (got ndim={mixed_qkv.ndim}).")
+    if mixed_qkv.stride(-1) != 1:
+        raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
+    if b.ndim != 2:
+        raise ValueError(f"`b` must be 2D (got b.ndim={b.ndim}).")
+    if A_log.ndim != 1:
+        raise ValueError("`A_log` must be a 1D tensor.")
+    if initial_state.ndim != 4:
+        raise ValueError(f"`initial_state` must be 4D (got ndim={initial_state.ndim}).")
+    if not out.is_contiguous():
+        raise ValueError("`out` must be contiguous.")
+    if write_pos.ndim != 1 or write_pos.dtype != torch.int32:
+        raise ValueError("`write_pos` must be a 1D int32 tensor.")
+    if force_flush is not None and (
+        force_flush.ndim != 1 or force_flush.dtype != torch.int32
+    ):
+        raise ValueError("`force_flush` must be a 1D int32 tensor or None.")
+
+    B = mixed_qkv.shape[0]
+    num_state_slots, HV, V, K = initial_state.shape
+    qkv_dim = mixed_qkv.shape[1]
+    q_dim = (qkv_dim - HV * V) // 2
+    if q_dim <= 0 or q_dim % K != 0:
+        raise ValueError(
+            f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}, K={K}."
+        )
+    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}."
+        )
+    max_cache_len = d_cache.shape[2]
+
+    # Gate-shape sanity: GDN scalar gate vs KDA per-K gate.
+    if is_kda:
+        if a.ndim != 3 or tuple(a.shape) != (B, HV, K):
+            raise ValueError(
+                f"KDA `a` must have shape {(B, HV, K)} (got {tuple(a.shape)})."
+            )
+        if dt_bias.ndim != 2 or tuple(dt_bias.shape) != (HV, K):
+            raise ValueError(
+                f"KDA `dt_bias` must have shape {(HV, K)} (got {tuple(dt_bias.shape)})."
+            )
+        if not a.is_contiguous() or not dt_bias.is_contiguous():
+            raise ValueError("KDA `a`/`dt_bias` must be contiguous.")
+        g_expect = (HV, max_cache_len, K)
+    else:
+        if a.ndim != 2 or tuple(a.shape) != (B, HV):
+            raise ValueError(
+                f"GDN `a` must have shape {(B, HV)} (got {tuple(a.shape)})."
+            )
+        if dt_bias.ndim != 1 or dt_bias.shape[0] != HV:
+            raise ValueError(
+                f"GDN `dt_bias` must have shape {(HV,)} (got {tuple(dt_bias.shape)})."
+            )
+        g_expect = (HV, max_cache_len)
+
+    # Cache shape sanity (per state slot): d=(HV, L, V), k=(H, L, K).
+    if tuple(d_cache.shape[1:]) != (HV, max_cache_len, V):
+        raise ValueError(
+            f"`d_cache` per-slot shape must be {(HV, max_cache_len, V)} "
+            f"(got {tuple(d_cache.shape[1:])})."
+        )
+    if tuple(k_cache.shape[1:]) != (H, max_cache_len, K):
+        raise ValueError(
+            f"`k_cache` per-slot shape must be {(H, max_cache_len, K)} "
+            f"(got {tuple(k_cache.shape[1:])})."
+        )
+    if tuple(g_cache.shape[1:]) != g_expect:
+        raise ValueError(
+            f"`g_cache` per-slot shape must be {g_expect} "
+            f"(got {tuple(g_cache.shape[1:])})."
+        )
+    if g_cache.dtype != torch.float32:
+        raise ValueError(f"`g_cache` must be float32 (got {g_cache.dtype}).")
+    if out.shape != (B, 1, HV, V):
+        raise ValueError(
+            f"`out` must have shape {(B, 1, HV, V)} (got {tuple(out.shape)})."
+        )
+    if write_pos.shape[0] != B or ssm_state_indices.shape[0] != B:
+        raise ValueError(
+            "`write_pos` and `ssm_state_indices` must both have length B="
+            f"{B} (got {write_pos.shape[0]}, {ssm_state_indices.shape[0]})."
+        )
+
+    BK = triton.next_power_of_2(K)
+    if triton.cdiv(K, BK) != 1:
+        raise ValueError(
+            f"Cached decode kernel only supports NK_global=1 (got K={K}, BK={BK})."
+        )
+    if BK % nk != 0:
+        raise ValueError(f"nk={nk} must divide BK={BK}.")
+    BKT = BK // nk
+    if BKT < 16:
+        raise ValueError(f"BKT={BKT} must be >=16 for tl.dot (nk={nk}, BK={BK}).")
+    # K-tiling keeps the per-program tile small enough that a larger BV fits
+    # without register spilling -> NV=1 -> half the grid -> fewer redundant
+    # cache / metadata loads.
+    BV = block_v if block_v is not None else min(triton.next_power_of_2(V), 64)
+    BC = max(16, triton.next_power_of_2(max_cache_len))
+
+    grid = (triton.cdiv(V, BV), B, HV)
+    fused_recurrent_linear_replayssm_decode_kernel[grid](
+        mixed_qkv=mixed_qkv,
+        a=a,
+        b=b,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        o=out,
+        h0=initial_state,
+        ht=initial_state,
+        d_cache=d_cache,
+        k_cache=k_cache,
+        g_cache=g_cache,
+        ssm_state_indices=ssm_state_indices,
+        write_pos=write_pos,
+        force_flush=force_flush if force_flush is not None else write_pos,
+        scale=scale,
+        stride_mixed_qkv_tok=mixed_qkv.stride(0),
+        stride_a_tok=a.stride(0),
+        stride_b_tok=b.stride(0),
+        stride_init_state_token=initial_state.stride(0),
+        stride_final_state_token=initial_state.stride(0),
+        stride_indices_seq=ssm_state_indices.stride(0),
+        H=H,
+        HV=HV,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        BC=BC,
+        NK=nk,
+        BKT=BKT,
+        MAX_CACHE_LEN=max_cache_len,
+        SOFTPLUS_THRESHOLD=20.0,
+        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        HAS_FORCE_FLUSH=force_flush is not None,
+        IS_KDA=is_kda,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return out, initial_state
+
+
+# Backwards-compatible aliases: the original GDN-only names. Existing callers
+# (backend dispatch, tests, microbench) keep working; ``is_kda`` defaults to
+# False so these are the GDN path unchanged.
+fused_recurrent_gdn_replayssm_decode_kernel = (
+    fused_recurrent_linear_replayssm_decode_kernel
+)
+fused_recurrent_gdn_replayssm_decode = fused_recurrent_linear_replayssm_decode

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

@@ -37,6 +37,14 @@ class MambaAttnBackendBase(AttentionBackend):
         self.token_to_kv_pool = model_runner.token_to_kv_pool
         self.forward_metadata: ForwardMetadata = None
         self.state_indices_list = []
+        # GDN ReplaySSM (slice 1b): per-bs STATIC per-row write-cursor buffers
+        # for cuda-graph. Allocated lazily in init_cuda_graph_state only when
+        # --enable-linear-replayssm is set; stays None otherwise.
+        self.replayssm_write_pos_list = None
+        # GDN ReplaySSM (slice 2b): per-bs STATIC per-row force-flush buffers
+        # for cuda-graph, parallel to replayssm_write_pos_list. Same lifetime
+        # (None unless the flag is on).
+        self.replayssm_force_flush_list = None
         self.query_start_loc_list = []
         self.retrieve_next_token_list = []
         self.retrieve_next_sibling_list = []
@@ -103,10 +111,83 @@ class MambaAttnBackendBase(AttentionBackend):
             mamba_cache_indices = mamba_cache_indices.clone()
             mamba_cache_indices[_real_bs:] = -1
 
+        replayssm_write_pos = None
+        replayssm_force_flush = None
         if forward_batch.forward_mode.is_decode_or_idle():
             query_start_loc = torch.arange(
                 0, bs + 1, dtype=torch.int32, device=self.device
             )
+            # GDN ReplaySSM (slice 1a): the ring cursor is a per-slot
+            # decode-position counter shared by ALL GDN layers in this forward.
+            # Manage it exactly ONCE here (not per-layer): snapshot this step's
+            # value for the batch's slots, hand it to the layers, then advance
+            # the persistent buffer mod L for the next step.
+            mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
+            write_pos_buf = (
+                getattr(mamba_pool, "replayssm_write_pos", None)
+                if mamba_pool is not None
+                else None
+            )
+            if write_pos_buf is not None:
+                slots = mamba_cache_indices.to(torch.long)
+                # Padded rows carry slot == -1; clamp so the per-row gather stays
+                # in-bounds (the kernel zeroes padded rows via state_idx < 0).
+                safe_slots = slots.clamp(min=0)
+                replayssm_write_pos = write_pos_buf[safe_slots].clone()
+                L = mamba_pool.linear_replayssm_cache_len
+                # KDA (per-K gate) ships without radix coordination for now: no
+                # track-boundary force-flush, so the ring flushes only at the
+                # natural write_pos == L-1 wrap. GDN keeps the radix-aligned
+                # force-flush (slice 2b). Gate on the pool's recorded gate type.
+                is_kda = getattr(mamba_pool, "replayssm_is_kda", False)
+                # GDN ReplaySSM (slice 2b): per-row force-flush at the radix
+                # track boundary. THE alignment: the radix mamba track snapshots
+                # temporal[slot] when seq_lens_cpu % mamba_track_interval == 0
+                # (extra_buffer: schedule_batch.prepare_for_decode builds
+                # `mamba_track_mask = (seq_lens_cpu % mamba_track_interval == 0)`
+                # off the SAME post-increment seq_lens_cpu used here). We source
+                # the flush from the identical seq_lens + condition so the kernel
+                # folds the ring into temporal[slot] on EXACTLY the steps the
+                # snapshot reads it. seq_lens_cpu is the committed length AFTER
+                # this decode token (incremented in prepare_for_decode before the
+                # forward), matching the track. int32, one entry per batch row.
+                if not is_kda:
+                    force_flush_bool = self._replayssm_track_flush_mask(
+                        forward_batch.seq_lens_cpu, bs
+                    )
+                    replayssm_force_flush = force_flush_bool.to(
+                        device=self.device, dtype=torch.int32
+                    )
+                # Advance only the VALID (non-padded) slots. Scatter over the
+                # unique valid slots to avoid duplicate-index races (padded rows
+                # all clamp to slot 0, which a real row may also occupy). A
+                # forced flush empties the ring -> next write_pos is 0 (same as
+                # the natural wrap at write_pos == L-1).
+                valid_mask = slots >= 0
+                valid_slots = slots[valid_mask]
+                if valid_slots.numel() > 0:
+                    # Per-row "did this step flush?": natural wrap OR forced.
+                    # (KDA has no forced flush -> force_flush is None -> pure wrap.)
+                    flushed = replayssm_write_pos == (L - 1)
+                    if replayssm_force_flush is not None:
+                        flushed = flushed | (replayssm_force_flush != 0)
+                    next_pos = torch.where(
+                        flushed,
+                        torch.zeros_like(replayssm_write_pos),
+                        (replayssm_write_pos + 1) % L,
+                    )
+                    # Dedup valid slots; for duplicates a scatter picks one
+                    # arbitrary row, but all rows of a given slot share the same
+                    # write_pos/flush, so the value is identical regardless.
+                    uniq_slots, inv = torch.unique(valid_slots, return_inverse=True)
+                    next_for_valid = next_pos[valid_mask]
+                    new_vals = torch.empty(
+                        uniq_slots.shape[0],
+                        dtype=write_pos_buf.dtype,
+                        device=write_pos_buf.device,
+                    )
+                    new_vals[inv] = next_for_valid.to(write_pos_buf.dtype)
+                    write_pos_buf[uniq_slots] = new_vals
         elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
             if forward_batch.forward_mode.is_draft_extend_v2():
                 # HybridLinearAttnBackend.init_forward_metadata calls all sub-backends
@@ -173,6 +254,8 @@ class MambaAttnBackendBase(AttentionBackend):
             track_ssm_final_src=track_ssm_final_src,
             track_ssm_final_dst=track_ssm_final_dst,
             has_mamba_track_mask=has_mamba_track_mask,
+            replayssm_write_pos=replayssm_write_pos,
+            replayssm_force_flush=replayssm_force_flush,
         )
 
     def init_forward_metadata_out_graph(
@@ -191,6 +274,7 @@ class MambaAttnBackendBase(AttentionBackend):
             num_padding=(
                 0 if in_capture else getattr(forward_batch, "num_padding", None)
             ),
+            in_capture=in_capture,
         )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
@@ -342,17 +426,72 @@ class MambaAttnBackendBase(AttentionBackend):
             bs, req_pool_indices, forward_mode, spec_info
         )
 
+    def _replayssm_enabled(self) -> bool:
+        """True iff --enable-linear-replayssm allocated the persistent ring cursor.
+
+        The per-slot ``replayssm_write_pos`` buffer on MambaPool is None unless
+        the flag is set, so it doubles as the on/off gate (same signal that
+        ``_forward_metadata`` / ``GDNAttnBackend.forward_decode`` already use).
+        """
+        mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
+        if mamba_pool is None:
+            return False
+        return getattr(mamba_pool, "replayssm_write_pos", None) is not None
+
+    def _replayssm_track_flush_mask(
+        self, seq_lens_cpu: torch.Tensor, bs: int
+    ) -> torch.Tensor:
+        """Per-row bool flush mask == the radix mamba-track snapshot condition.
+
+        THE alignment (slice 2b): the radix mamba track snapshots temporal[slot]
+        exactly when ``seq_lens_cpu % mamba_track_interval == 0`` (the same mask
+        ``schedule_batch.prepare_for_decode`` builds for extra_buffer at
+        ``mamba_track_mask = (seq_lens_cpu % mamba_track_interval == 0)``). Both
+        read the SAME post-increment ``seq_lens_cpu`` (committed length AFTER
+        this decode token), so the kernel force-flush fires on EXACTLY the steps
+        the snapshot reads the checkpoint -- no off-by-one. Returns a CPU bool
+        tensor of length ``bs`` (caller moves it to device as int32).
+        """
+        interval = get_global_server_args().mamba_track_interval
+        if seq_lens_cpu is None:
+            # Decode without a CPU seq-len mirror should not happen for the
+            # supported (no_buffer, radix-on) config, but stay safe: never flush.
+            return torch.zeros((bs,), dtype=torch.bool)
+        mask = (seq_lens_cpu[:bs].to(torch.int64) % interval) == 0
+        if mask.shape[0] < bs:
+            pad = torch.zeros((bs - mask.shape[0],), dtype=torch.bool)
+            mask = torch.cat([mask, pad])
+        return mask.cpu()
+
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         assert (
             max_num_tokens % max_bs == 0
         ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}"
         draft_token_num = max_num_tokens // max_bs
+        # GDN ReplaySSM (slice 1b): per-batch-size STATIC per-row write-cursor
+        # buffers the kernel reads. Captured into the graph by pointer, so they
+        # must be the SAME tensor objects refreshed in-place each replay. Sized
+        # and indexed like state_indices_list ((i+1,), indexed [bs - 1]). Left
+        # None when the flag is off so the dispatch falls through unchanged.
+        self.replayssm_write_pos_list = [] if self._replayssm_enabled() else None
+        # GDN ReplaySSM (slice 2b): static per-bs force-flush buffers, captured
+        # by pointer and refreshed in-place per replay just like the write-pos
+        # buffers. None when the flag is off.
+        self.replayssm_force_flush_list = [] if self._replayssm_enabled() else None
         for i in range(max_bs):
             self.state_indices_list.append(
                 torch.full(
                     (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device
                 )
             )
+            if self.replayssm_write_pos_list is not None:
+                self.replayssm_write_pos_list.append(
+                    torch.zeros((i + 1,), dtype=torch.int32, device=self.device)
+                )
+            if self.replayssm_force_flush_list is not None:
+                self.replayssm_force_flush_list.append(
+                    torch.zeros((i + 1,), dtype=torch.int32, device=self.device)
+                )
             self.query_start_loc_list.append(
                 torch.zeros((i + 2,), dtype=torch.int32, device=self.device)
             )
@@ -419,6 +558,25 @@ class MambaAttnBackendBase(AttentionBackend):
         mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
         self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
 
+        # GDN ReplaySSM (slice 1b): point at the STATIC per-bs write-cursor
+        # buffer (no advance, no snapshot — capture records the pointer; its
+        # zeros are overwritten in-place by _replay_metadata before each
+        # replay). None when the flag is off. Same per-bs tensor object that
+        # _replay_metadata refreshes, so the captured pointer stays valid.
+        replayssm_write_pos = (
+            self.replayssm_write_pos_list[bs - 1]
+            if self.replayssm_write_pos_list is not None
+            else None
+        )
+        # GDN ReplaySSM (slice 2b): point at the STATIC per-bs force-flush
+        # buffer (same capture-by-pointer contract as write_pos; refreshed
+        # in-place by _replay_metadata before each replay). None when off.
+        replayssm_force_flush = (
+            self.replayssm_force_flush_list[bs - 1]
+            if self.replayssm_force_flush_list is not None
+            else None
+        )
+
         # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
         if forward_mode.is_target_verify() and self.topk > 1:
             # They are None during cuda graph capture so skip the copy_...
@@ -430,11 +588,15 @@ class MambaAttnBackendBase(AttentionBackend):
                 retrieve_next_token=self.retrieve_next_token_list[bs - 1],
                 retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1],
                 retrieve_parent_token=self.retrieve_parent_token_list[bs - 1],
+                replayssm_write_pos=replayssm_write_pos,
+                replayssm_force_flush=replayssm_force_flush,
             )
         else:
             return ForwardMetadata(
                 query_start_loc=self.query_start_loc_list[bs - 1],
                 mamba_cache_indices=self.state_indices_list[bs - 1],
+                replayssm_write_pos=replayssm_write_pos,
+                replayssm_force_flush=replayssm_force_flush,
             )
 
     def _replay_metadata(
@@ -445,6 +607,7 @@ class MambaAttnBackendBase(AttentionBackend):
         spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],
         num_padding: Optional[int] = None,
+        in_capture: bool = False,
     ):
         if num_padding is None:
             if seq_lens_cpu is None:
@@ -458,6 +621,84 @@ class MambaAttnBackendBase(AttentionBackend):
         mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
         mamba_indices[bs - num_padding :] = -1
         self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
+        # GDN ReplaySSM (slice 1b): refresh the STATIC per-row write cursor the
+        # kernel reads, mirroring the eager snapshot-then-advance in
+        # _forward_metadata but writing in-place into the captured per-bs buffer
+        # so the graph's recorded pointer stays valid across replays. Done once
+        # per forward here (out_graph host op), not per layer. Skipped during
+        # capture (in_capture): capture runs on dummy slots, so advancing the
+        # persistent counter then would corrupt real per-slot ring positions;
+        # the captured buffer's contents are irrelevant at capture time anyway.
+        replayssm_write_pos = None
+        replayssm_force_flush = None
+        if self.replayssm_write_pos_list is not None:
+            mamba_pool = self.req_to_token_pool.mamba_pool
+            write_pos_buf = mamba_pool.replayssm_write_pos
+            static_wp = self.replayssm_write_pos_list[bs - 1]
+            static_ff = self.replayssm_force_flush_list[bs - 1]
+            # Hand the full captured per-bs buffers to the kernel, mirroring how
+            # mamba_cache_indices = self.state_indices_list[bs - 1] is the full
+            # (bs,) tensor; the kernel indexes them per decode row.
+            replayssm_write_pos = static_wp
+            replayssm_force_flush = static_ff
+            if write_pos_buf is not None:
+                # mamba_indices: this replay's per-row physical slots (padded
+                # rows == -1, same tensor fed to state_indices_list above).
+                slots = mamba_indices.to(torch.long)
+                safe_slots = slots.clamp(min=0)
+                # Snapshot THIS step's per-slot cursor into the captured buffer
+                # the kernel reads (in-place copy_, never reassign the object).
+                static_wp[: len(mamba_indices)].copy_(write_pos_buf[safe_slots])
+                # GDN ReplaySSM (slice 2b): refresh the captured force-flush
+                # buffer in-place from THIS step's seq_lens. THE alignment: same
+                # `seq_lens_cpu % mamba_track_interval == 0` the radix track uses
+                # (see _replayssm_track_flush_mask / schedule_batch). During
+                # capture (seq_lens_cpu is None) leave it zeroed: capture content
+                # is irrelevant and decode replays overwrite it below.
+                force_flush_dev = None
+                # KDA: no radix coordination -> leave static_ff zeroed and
+                # force_flush_dev None so the advance below is a pure wrap,
+                # matching the kernel (a zeroed force_flush flushes nothing).
+                is_kda = getattr(mamba_pool, "replayssm_is_kda", False)
+                if (
+                    not is_kda
+                    and forward_mode.is_decode_or_idle()
+                    and seq_lens_cpu is not None
+                ):
+                    ff_mask = self._replayssm_track_flush_mask(seq_lens_cpu, bs)
+                    force_flush_dev = ff_mask.to(device=self.device, dtype=torch.int32)
+                    static_ff.copy_(force_flush_dev)
+                else:
+                    static_ff.zero_()
+                if not in_capture:
+                    L = mamba_pool.linear_replayssm_cache_len
+                    # Advance only VALID (non-padded) slots. A forced flush
+                    # empties the ring -> next write_pos is 0 (same as the
+                    # natural wrap at write_pos == L-1). Use this step's snapshot
+                    # cursor (write_pos_buf[safe_slots]) + the flush flag.
+                    valid_mask = slots >= 0
+                    valid_slots = slots[valid_mask]
+                    if valid_slots.numel() > 0:
+                        cur_pos = write_pos_buf[safe_slots]
+                        flushed = cur_pos == (L - 1)
+                        if force_flush_dev is not None:
+                            flushed = flushed | (force_flush_dev != 0)
+                        next_pos = torch.where(
+                            flushed,
+                            torch.zeros_like(cur_pos),
+                            (cur_pos + 1) % L,
+                        )
+                        # Dedup; rows sharing a slot share write_pos+flush, so
+                        # the scattered value is identical for either row.
+                        uniq_slots, inv = torch.unique(valid_slots, return_inverse=True)
+                        next_for_valid = next_pos[valid_mask]
+                        new_vals = torch.empty(
+                            uniq_slots.shape[0],
+                            dtype=write_pos_buf.dtype,
+                            device=write_pos_buf.device,
+                        )
+                        new_vals[inv] = next_for_valid.to(write_pos_buf.dtype)
+                        write_pos_buf[uniq_slots] = new_vals
         if forward_mode.is_decode_or_idle():
             if num_padding == 0:
                 self.query_start_loc_list[bs - 1].copy_(
@@ -504,11 +745,15 @@ class MambaAttnBackendBase(AttentionBackend):
                 retrieve_next_token=self.retrieve_next_token_list[bs - 1],
                 retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1],
                 retrieve_parent_token=self.retrieve_parent_token_list[bs - 1],
+                replayssm_write_pos=replayssm_write_pos,
+                replayssm_force_flush=replayssm_force_flush,
             )
         else:
             return ForwardMetadata(
                 query_start_loc=self.query_start_loc_list[bs - 1],
                 mamba_cache_indices=self.state_indices_list[bs - 1],
+                replayssm_write_pos=replayssm_write_pos,
+                replayssm_force_flush=replayssm_force_flush,
             )
 
     def get_cuda_graph_seq_len_fill_value(self):
@@ -615,6 +860,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
             num_padding=(
                 0 if in_capture else getattr(forward_batch, "num_padding", None)
             ),
+            in_capture=in_capture,
         )
         spec_info = forward_batch.spec_info
         draft_token_num = spec_info.draft_token_num if spec_info is not None else 1

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

@@ -314,6 +314,17 @@ class GDNAttnBackend(MambaAttnBackendBase):
         ssm_states = layer_cache.temporal
         query_start_loc = self.forward_metadata.query_start_loc
         cache_indices = self.forward_metadata.mamba_cache_indices
+        # GDN ReplaySSM (slice 1a): per-layer ring slices + the once-per-forward
+        # per-row write cursor. All None unless --enable-linear-replayssm, so the
+        # legacy dispatch below is byte-identical when the flag is off.
+        replayssm_write_pos = self.forward_metadata.replayssm_write_pos
+        # GDN ReplaySSM (slice 2b): per-row force-flush at radix track
+        # boundaries (None unless --enable-linear-replayssm). When present the
+        # kernel folds the ring into temporal[slot] on the snapshot steps.
+        replayssm_force_flush = self.forward_metadata.replayssm_force_flush
+        replayssm_d = layer_cache.replayssm_d
+        replayssm_k = layer_cache.replayssm_k
+        replayssm_g = layer_cache.replayssm_g
 
         assert isinstance(mixed_qkv, torch.Tensor)
         mixed_qkv = causal_conv1d_update(
@@ -339,6 +350,11 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 cache_indices=cache_indices,
                 num_v_heads=layer.num_v_heads,
                 head_v_dim=layer.head_v_dim,
+                replayssm_d=replayssm_d,
+                replayssm_k=replayssm_k,
+                replayssm_g=replayssm_g,
+                replayssm_write_pos=replayssm_write_pos,
+                replayssm_force_flush=replayssm_force_flush,
             )
             self._track_mamba_state_decode(
                 forward_batch, conv_states, ssm_states, cache_indices

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

@@ -205,6 +205,28 @@ class KDAAttnBackend(MambaAttnBackendBase):
         query_start_loc = self.forward_metadata.query_start_loc
         cache_indices = self.forward_metadata.mamba_cache_indices
 
+        # ReplaySSM ring: per-layer ring slices + the once-per-forward per-row
+        # write cursor. All None unless --enable-linear-replayssm, so packed_decode
+        # falls through to the byte-identical legacy KDA path. KDA ships WITHOUT
+        # radix coordination for now, so force_flush is None/zeroed (the ring
+        # flushes only at the natural write_pos == L-1 wrap; set in the shared
+        # HybridLinearAttn metadata, which zeroes force_flush for KDA models).
+        # NOTE: ReplaySSM decode is a GDN (scalar-gate) bandwidth win; on KDA the
+        # per-K g_cache is K x larger and the reconstruction refolds the per-K
+        # decay every step, so it is correct but SLOWER than packed (a measured
+        # decode regression). Kept wired for correctness + the spec-decode path;
+        # not recommended for KDA decode. Revisit on Blackwell (more tensor-core
+        # throughput may flip the compute/bandwidth tradeoff).
+        replayssm_write_pos = getattr(
+            self.forward_metadata, "replayssm_write_pos", None
+        )
+        replayssm_force_flush = getattr(
+            self.forward_metadata, "replayssm_force_flush", None
+        )
+        replayssm_d = layer_cache.replayssm_d
+        replayssm_k = layer_cache.replayssm_k
+        replayssm_g = layer_cache.replayssm_g
+
         qkv = causal_conv1d_update(
             mixed_qkv,
             conv_states.transpose(-1, -2),
@@ -240,6 +262,11 @@ class KDAAttnBackend(MambaAttnBackendBase):
                 cache_indices=cache_indices,
                 num_v_heads=layer.num_v_heads,
                 head_v_dim=layer.head_v_dim,
+                replayssm_d=replayssm_d,
+                replayssm_k=replayssm_k,
+                replayssm_g=replayssm_g,
+                replayssm_write_pos=replayssm_write_pos,
+                replayssm_force_flush=replayssm_force_flush,
             )
 
         q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)

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

@@ -10,6 +10,9 @@ if not is_cpu():
     from sglang.srt.layers.attention.fla.fused_recurrent import (
         fused_recurrent_gated_delta_rule_packed_decode,
     )
+    from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+        fused_recurrent_gdn_replayssm_decode,
+    )
     from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
         fused_sigmoid_gating_delta_rule_update,
     )
@@ -78,6 +81,43 @@ class TritonGDNKernel(LinearAttnKernelBase):
         # Packed kernel expects output shape [B, 1, HV, V]
         out = mixed_qkv.new_empty(B, 1, num_v_heads, head_v_dim)
 
+        # GDN ReplaySSM buffered decode (slice 1a). Drop-in for the packed
+        # decode: same args plus the three per-layer ring caches and the
+        # per-row write cursor. When any ring tensor / cursor is None (flag
+        # off) we fall through to the byte-identical legacy path below.
+        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")
+        # GDN ReplaySSM (slice 2b): optional per-row force-flush (radix track
+        # boundary). None when radix tracking is off / flag off; the kernel
+        # treats None as "no forced flush" (byte-identical to slice 1a/1b).
+        replayssm_force_flush = kwargs.get("replayssm_force_flush")
+        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
+        ):
+            fused_recurrent_gdn_replayssm_decode(
+                mixed_qkv=mixed_qkv,
+                a=a,
+                b=b,
+                A_log=A_log,
+                dt_bias=dt_bias,
+                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=replayssm_force_flush,
+                use_qk_l2norm_in_kernel=True,
+            )
+            return out.transpose(0, 1)
+
         fused_recurrent_gated_delta_rule_packed_decode(
             mixed_qkv=mixed_qkv,
             a=a,

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

@@ -11,6 +11,9 @@ if not is_cpu():
     from sglang.srt.layers.attention.fla.fused_recurrent import (
         fused_recurrent_kda_packed_decode,
     )
+    from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+        fused_recurrent_linear_replayssm_decode,
+    )
     from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
         fused_sigmoid_gating_delta_rule_update,
     )
@@ -45,13 +48,52 @@ class TritonKDAKernel(LinearAttnKernelBase):
         decode kernel output layout.
         """
         B = mixed_qkv.shape[0]
+        out = mixed_qkv.new_empty(B, 1, num_v_heads, head_v_dim)
+
+        # KDA ReplaySSM buffered decode: drop-in for the packed decode, same
+        # args plus the three per-layer ring caches + the per-row write cursor
+        # (and optional radix-track force-flush). Uses the gate-generic kernel
+        # with is_kda=True (per-K gate); g_cache is [num_slots, HV, L, K].
+        # When any ring tensor / cursor is None (flag off) we fall through to
+        # the byte-identical legacy path below.
+        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")
+        replayssm_force_flush = kwargs.get("replayssm_force_flush")
+        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
+        ):
+            K = ssm_states.shape[-1]  # ssm_states: [num_slots, HV, V, K]
+            fused_recurrent_linear_replayssm_decode(
+                mixed_qkv=mixed_qkv,
+                a=a.reshape(B, num_v_heads, K).contiguous(),
+                b=b.reshape(B, num_v_heads).contiguous(),
+                A_log=A_log.reshape(-1),
+                dt_bias=dt_bias.reshape(num_v_heads, K).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=replayssm_force_flush,
+                use_qk_l2norm_in_kernel=True,
+                is_kda=True,
+            )
+            return out.transpose(0, 1)
+
         # a may come in as [B, HV, K] (or [B, 1, HV*K]); b may come in as
         # [B, 1, HV]. Flatten both to the 2D shapes the kernel expects.
         if a.dim() != 2:
             a = a.reshape(B, -1)
         if b.dim() != 2:
             b = b.reshape(B, -1)
-        out = mixed_qkv.new_empty(B, 1, num_v_heads, head_v_dim)
         fused_recurrent_kda_packed_decode(
             mixed_qkv=mixed_qkv,
             a=a,

python/sglang/srt/layers/attention/mamba/mamba2_metadata.py

@@ -30,6 +30,16 @@ class ForwardMetadata:
     query_start_loc: torch.Tensor
     mamba_cache_indices: torch.Tensor
     mamba_cache_indices_gdn: Optional[torch.Tensor] = None
+    # GDN ReplaySSM (slice 1a): per-decode-row snapshot of the ring write
+    # cursor for THIS decode step (gathered from the persistent per-slot
+    # buffer, then advanced once for the next step). int32, length == batch.
+    replayssm_write_pos: Optional[torch.Tensor] = None
+    # GDN ReplaySSM (slice 2b): per-decode-row int32 flush flag for THIS decode
+    # step. !=0 forces the kernel to fold the partial ring + current token into
+    # the checkpoint (temporal[slot]) so the radix cache reads an up-to-date
+    # state. Fires on EXACTLY the rows the radix track snapshots, i.e. the same
+    # condition the track uses: seq_lens_cpu % mamba_track_interval == 0.
+    replayssm_force_flush: Optional[torch.Tensor] = None
     # For topk > 1 eagle
     retrieve_next_token: Optional[torch.Tensor] = None
     retrieve_next_sibling: Optional[torch.Tensor] = None

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -530,11 +530,19 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         ]
 
         if is_insert:
-            cache_len = (
+            if self.enable_mamba_extra_buffer:
-                req.mamba_last_track_seqlen
+                cache_len = req.mamba_last_track_seqlen
-                if self.enable_mamba_extra_buffer
+            else:
-                else len(token_ids)
+                cache_len = len(token_ids)
-            )
+                # ReplaySSM (no_buffer): `temporal[slot]` lags the live state by
+                # the slot's unflushed ring depth (`write_pos`), so cap the
+                # donate to the last flush boundary (where temporal is current)
+                # and reset the cursor, keeping the donated checkpoint consistent
+                # with its key length. page_size is asserted == 1, so no realign.
+                write_pos_buf = self.req_to_token_pool.mamba_pool.replayssm_write_pos
+                if write_pos_buf is not None:
+                    cache_len -= int(write_pos_buf[req.mamba_pool_idx].item())
+                    write_pos_buf[req.mamba_pool_idx] = 0
             if cache_len is None:
                 cache_len = 0
             if cache_len != len(token_ids):

python/sglang/srt/mem_cache/memory_pool.py

@@ -307,6 +307,15 @@ class MambaPool:
     class State:
         conv: List[torch.Tensor]
         temporal: torch.Tensor
+        # GDN ReplaySSM ring buffers (slice 1a). Only allocated when
+        # `--enable-linear-replayssm` is set; otherwise None so the legacy path is
+        # byte-identical. Per-layer layout: [num_layers, num_slots, ...].
+        #   replayssm_d: [num_layers, num_slots, HV, L, V]
+        #   replayssm_k: [num_layers, num_slots, H,  L, K]
+        #   replayssm_g: [num_layers, num_slots, HV, L]  (fp32)
+        replayssm_d: Optional[torch.Tensor] = None
+        replayssm_k: Optional[torch.Tensor] = None
+        replayssm_g: Optional[torch.Tensor] = None
 
         def at_layer_idx(self, layer: int):
             kwargs = {}
@@ -314,7 +323,9 @@ class MambaPool:
             for f in fields(self):
                 name = f.name
                 v = getattr(self, name)
-                if name in ("conv", "intermediate_conv_window"):
+                if v is None:
+                    kwargs[name] = None
+                elif name in ("conv", "intermediate_conv_window"):
                     kwargs[name] = [conv[layer] for conv in v]
                 else:
                     kwargs[name] = v[layer]
@@ -325,6 +336,7 @@ class MambaPool:
             return sum(
                 get_tensor_size_bytes(getattr(self, f.name))
                 for f in dataclasses.fields(self)
+                if getattr(self, f.name) is not None
             )
 
     @dataclass(frozen=True, kw_only=True)
@@ -343,6 +355,8 @@ class MambaPool:
         enable_memory_saver: bool = False,
         speculative_num_draft_tokens: Optional[int] = None,
         speculative_eagle_topk: Optional[int] = None,
+        enable_linear_replayssm: bool = False,
+        linear_replayssm_cache_len: int = 16,
     ):
         conv_state_shape = cache_params.shape.conv
         temporal_state_shape = cache_params.shape.temporal
@@ -355,6 +369,8 @@ class MambaPool:
 
         self.size = size
         self.device = device
+        self.enable_linear_replayssm = enable_linear_replayssm
+        self.linear_replayssm_cache_len = linear_replayssm_cache_len
 
         # for disagg with nvlink
         self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
@@ -398,6 +414,41 @@ class MambaPool:
                 dtype=ssm_dtype,
                 device=device,
             )
+
+            # GDN ReplaySSM ring buffers (slice 1a). Allocated only when the
+            # flag is on; otherwise left as None so the legacy State is
+            # byte-identical. temporal_state_shape == (HV, V, K).
+            replayssm_d = replayssm_k = replayssm_g = None
+            if enable_linear_replayssm:
+                hv, v_dim, k_dim = temporal_state_shape
+                h_k = getattr(cache_params.shape, "num_k_heads_per_tp", hv)
+                L = linear_replayssm_cache_len
+                num_slots = size + 1
+                # Ring records live in the SSM dtype (bf16/fp32) except g (fp32).
+                replayssm_d = torch.zeros(
+                    size=(num_mamba_layers, num_slots, hv, L, v_dim),
+                    dtype=ssm_dtype,
+                    device=device,
+                )
+                replayssm_k = torch.zeros(
+                    size=(num_mamba_layers, num_slots, h_k, L, k_dim),
+                    dtype=ssm_dtype,
+                    device=device,
+                )
+                # The log-decay gate ring (fp32): per-head SCALAR for the GDN
+                # gate -> [.., L]; per-K VECTOR for the KDA gate -> [.., L, K]
+                # (k_dim == temporal_state_shape[-1] for both).
+                g_shape = (
+                    (num_mamba_layers, num_slots, hv, L, k_dim)
+                    if cache_params.is_kda
+                    else (num_mamba_layers, num_slots, hv, L)
+                )
+                replayssm_g = torch.zeros(
+                    size=g_shape,
+                    dtype=torch.float32,
+                    device=device,
+                )
+
             if speculative_num_draft_tokens is not None:
                 if _is_npu:
                     temporal_state = temporal_state.transpose(-1, -2)
@@ -505,6 +556,9 @@ class MambaPool:
                     temporal=temporal_state,
                     intermediate_ssm=intermediate_ssm_state_cache,
                     intermediate_conv_window=intermediate_conv_window_cache,
+                    replayssm_d=replayssm_d,
+                    replayssm_k=replayssm_k,
+                    replayssm_g=replayssm_g,
                 )
                 logger.info(
                     f"Mamba Cache is allocated. "
@@ -517,13 +571,41 @@ class MambaPool:
                     f"intermediate_conv_window_cache size: {get_tensor_size_bytes(self._intermediate_conv_window_phys) / GB:.2f}GB "
                 )
             else:
-                self.mamba_cache = self.State(conv=conv_state, temporal=temporal_state)
+                self.mamba_cache = self.State(
+                    conv=conv_state,
+                    temporal=temporal_state,
+                    replayssm_d=replayssm_d,
+                    replayssm_k=replayssm_k,
+                    replayssm_g=replayssm_g,
+                )
                 logger.info(
                     f"Mamba Cache is allocated. "
                     f"max_mamba_cache_size: {size}, "
                     f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
                     f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
                 )
+            if enable_linear_replayssm:
+                logger.info(
+                    f"GDN ReplaySSM ring buffers allocated (L="
+                    f"{linear_replayssm_cache_len}): "
+                    f"d={get_tensor_size_bytes(replayssm_d) / GB:.3f}GB, "
+                    f"k={get_tensor_size_bytes(replayssm_k) / GB:.3f}GB, "
+                    f"g={get_tensor_size_bytes(replayssm_g) / GB:.3f}GB "
+                )
+            # Gate granularity of the linear-attn layers (drives the kernel's
+            # IS_KDA path + the g_cache layout). Read by the backend metadata to
+            # decide the per-K (KDA) vs scalar (GDN) flush/advance handling.
+            self.replayssm_is_kda = bool(
+                enable_linear_replayssm and cache_params.is_kda
+            )
+            # Persistent per-slot decode-position cursor for ReplaySSM. Shared
+            # across all linear-attn layers; advanced once per decode forward by
+            # the backend metadata build. Index 0..size; reset on slot (re)alloc.
+            self.replayssm_write_pos = (
+                torch.zeros((size + 1,), dtype=torch.int32, device=device)
+                if enable_linear_replayssm
+                else None
+            )
             mem_usage_bytes = self.mamba_cache.mem_usage_bytes()
             if isinstance(self.mamba_cache, self.SpeculativeState):
                 # `intermediate_conv_window` is an as_strided view whose logical
@@ -562,6 +644,24 @@ class MambaPool:
         t[:, indices] = z
 
     def copy_from(self, src_indices: torch.Tensor, dst_indices: torch.Tensor):
+        """Clone mamba state (conv + temporal) from src slots into dst slots.
+
+        ReplaySSM invariant: the SOURCE must be a fully-flushed checkpoint
+        (``write_pos[src] == 0``). Only ``temporal`` is copied, not the ring, so
+        an un-flushed source would drop its last ``write_pos`` updates. Callers
+        comply: COW copies radix checkpoints; ``cache_unfinished_req`` copies an
+        active slot only during prefill (ring empty); ``cache_finished_req``
+        caps the donate to the last flush boundary. The dst cursor is reset to 0
+        (the copied checkpoint has no pending ring entries).
+        """
+        if self.replayssm_write_pos is not None and envs.SGLANG_DEBUG_MEMORY_POOL.get():
+            # Debug-only (syncs): catch any copy of an active, un-flushed slot.
+            src_wp = self.replayssm_write_pos[src_indices]
+            assert bool((src_wp == 0).all().item()), (
+                "copy_from requires a fully-flushed ReplaySSM source "
+                f"(write_pos==0), got {src_wp.tolist()} for src "
+                f"{src_indices.tolist()}"
+            )
         for i in range(len(self.mamba_cache.conv)):
             self.mamba_cache.conv[i][:, dst_indices] = self.mamba_cache.conv[i][
                 :, src_indices
@@ -569,6 +669,8 @@ class MambaPool:
         self.mamba_cache.temporal[:, dst_indices] = self.mamba_cache.temporal[
             :, src_indices
         ]
+        if self.replayssm_write_pos is not None:
+            self.replayssm_write_pos[dst_indices] = 0
 
     def get_cpu_copy(self, indices):
         current_platform.synchronize()
@@ -604,7 +706,13 @@ class MambaPool:
             # These buffers have different size (spec_state_size + 1) and should not be transferred
             if field in ("intermediate_ssm", "intermediate_conv_window"):
                 continue
+            # Skip GDN ReplaySSM ring buffers: they are derived/transient decode
+            # scratch, not part of the persistent transferable state.
+            if field in ("replayssm_d", "replayssm_k", "replayssm_g"):
+                continue
             value = getattr(self.mamba_cache, field)
+            if value is None:
+                continue
             if isinstance(value, list):
                 state_tensors.extend(value)
             else:
@@ -633,7 +741,19 @@ class MambaPool:
         """
         state_tensors = []
         for field in vars(self.mamba_cache):
+            # Mirror the exclusions in get_contiguous_buf_infos so the returned
+            # dims line up element-wise with the RDMA buffer list.
+            if field in (
+                "intermediate_ssm",
+                "intermediate_conv_window",
+                "replayssm_d",
+                "replayssm_k",
+                "replayssm_g",
+            ):
+                continue
             value = getattr(self.mamba_cache, field)
+            if value is None:
+                continue
             if isinstance(value, list):
                 state_tensors.extend(value)
             else:
@@ -669,6 +789,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
         speculative_eagle_topk: Optional[int] = None,
         enable_overlap_schedule: bool = True,
         start_layer: Optional[int] = None,
+        enable_linear_replayssm: bool = False,
+        linear_replayssm_cache_len: int = 16,
     ):
         super().__init__(
             size=size,
@@ -692,6 +814,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
             enable_mamba_extra_buffer=enable_mamba_extra_buffer,
             speculative_num_draft_tokens=speculative_num_draft_tokens,
             speculative_eagle_topk=speculative_eagle_topk,
+            enable_linear_replayssm=enable_linear_replayssm,
+            linear_replayssm_cache_len=linear_replayssm_cache_len,
         )
 
     def _init_mamba_pool(
@@ -704,6 +828,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
         enable_mamba_extra_buffer: bool,
         speculative_num_draft_tokens: int = None,
         speculative_eagle_topk: Optional[int] = None,
+        enable_linear_replayssm: bool = False,
+        linear_replayssm_cache_len: int = 16,
     ):
         self.mamba_pool = MambaPool(
             size=mamba_size,
@@ -714,6 +840,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
             enable_memory_saver=self.enable_memory_saver,
             speculative_num_draft_tokens=speculative_num_draft_tokens,
             speculative_eagle_topk=speculative_eagle_topk,
+            enable_linear_replayssm=enable_linear_replayssm,
+            linear_replayssm_cache_len=linear_replayssm_cache_len,
         )
         self.mamba_allocator = MambaSlotAllocator(
             size=mamba_size,
@@ -771,6 +899,12 @@ class HybridReqToTokenPool(ReqToTokenPool):
                 ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_allocator.available_size()=}, {len(reqs)=}"
                 req.mamba_pool_idx = mid[0]
                 req.mamba_needs_clear = True
+                # GDN ReplaySSM: a freshly (re)assigned slot starts an empty
+                # ring. write_pos=0 means "ring empty", so the decode kernel
+                # ignores ring contents and reads only the checkpoint state
+                # (the post-prefill state that prefill wrote into this slot).
+                if self.mamba_pool.replayssm_write_pos is not None:
+                    self.mamba_pool.replayssm_write_pos[req.mamba_pool_idx] = 0
             mamba_indices.append(req.mamba_pool_idx)
             if self.enable_mamba_extra_buffer:
                 if req.mamba_ping_pong_track_buffer is None:

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -404,6 +404,8 @@ class ModelRunnerKVCacheMixin:
                     speculative_eagle_topk=self.server_args.speculative_eagle_topk,
                     enable_overlap_schedule=not self.server_args.disable_overlap_schedule,
                     start_layer=self.start_layer,
+                    enable_linear_replayssm=self.server_args.enable_linear_replayssm,
+                    linear_replayssm_cache_len=self.server_args.linear_replayssm_cache_len,
                 )
             else:
                 # DSV4 on NPU needs an extended ReqToTokenPool holding per-req

python/sglang/srt/server_args.py

@@ -1837,6 +1837,23 @@ class ServerArgs:
             choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
         ),
     ] = None
+    # ReplaySSM buffered output-only linear-attn decode (GDN + KDA): per-slot
+    # ring + periodic flush to cut per-step HBM state traffic.
+    enable_linear_replayssm: A[
+        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 "
+        "is numerically correct, but KDA decode is SLOWER than the packed "
+        "baseline (the per-K g_cache is K x larger and the reconstruction "
+        "refolds the per-K decay every step), so it is not recommended for KDA "
+        "models. Requires the Triton linear-attn decode backend and "
+        "--mamba-scheduler-strategy no_buffer (the default).",
+    ] = False
+    linear_replayssm_cache_len: A[
+        int,
+        "Ring-buffer length L for ReplaySSM linear-attn decode. The full recurrent state is flushed to HBM every L decode steps.",
+    ] = 16
 
     # -------------------------------------------------------------------------
     # Hierarchical cache
@@ -5052,6 +5069,50 @@ class ServerArgs:
                 f"got CUDA {cuda_version or 'unknown'}"
             )
 
+        # GDN ReplaySSM buffered decode guards. Runs on the Triton GDN decode
+        # backend. 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
+        # the COW copy-into-slot path resets the ring cursor) -- so the
+        # --disable-radix-cache requirement is dropped.
+        #
+        # Slice 2b only wires the no_buffer mamba scheduler strategy (the
+        # default). The extra_buffer strategy donates the track snapshot via
+        # `donate_mamba_ping_pong_slot` with a separate ping-pong slot swap that
+        # does NOT route through MambaPool.copy_from, so the ReplaySSM ring
+        # 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":
+                raise ValueError(
+                    "--enable-linear-replayssm requires the Triton "
+                    "linear-attn decode backend, got "
+                    f"--linear-attn-decode-backend={decode!r}."
+                )
+            if self.enable_mamba_extra_buffer():
+                raise ValueError(
+                    "--enable-linear-replayssm requires --mamba-scheduler-strategy "
+                    "no_buffer (the default); the extra_buffer ping-pong "
+                    "donation path is not yet supported (follow-up). Got "
+                    f"--mamba-scheduler-strategy={self.mamba_scheduler_strategy!r}."
+                )
+            if self.disaggregation_mode != "null":
+                # The disaggregated decode pool (HybridMambaDecodeReqToTokenPool)
+                # is not wired for the ReplaySSM ring, so the flag would silently
+                # no-op there; disagg also runs a different cache/coordination
+                # flow that is not yet validated for ReplaySSM (follow-up).
+                raise ValueError(
+                    "--enable-linear-replayssm is not supported under PD "
+                    "disaggregation yet (follow-up). Got "
+                    f"--disaggregation-mode={self.disaggregation_mode!r}."
+                )
+            if self.linear_replayssm_cache_len < 1:
+                raise ValueError(
+                    "--linear-replayssm-cache-len must be >= 1, got "
+                    f"{self.linear_replayssm_cache_len}."
+                )
+
     def _handle_legacy_cp_arguments(self):
         legacy_mode_to_strategy = {
             "in-seq-split": "zigzag",

test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py

@@ -0,0 +1,344 @@
+"""Correctness test for the buffered output-only linear-attention decode kernel
+(ReplaySSM Part A), implemented in
+``python/sglang/srt/layers/attention/fla/fused_recurrent_linear_replayssm.py``.
+
+One kernel, two gate granularities (selected by ``is_kda``):
+  * GDN (``is_kda=False``): per-head SCALAR gate. Ground truth =
+    ``fused_recurrent_gated_delta_rule_packed_decode`` (fused_recurrent.py:186).
+  * KDA (``is_kda=True``):  per-K-channel gate. Ground truth =
+    ``fused_recurrent_kda_packed_decode`` (fused_recurrent.py:406). The KDA
+    baseline takes the per-K gate flattened (``a``=[B, HV*K], ``dt_bias``=[HV*K]);
+    the ReplaySSM kernel takes it as [B, HV, K] / [HV, K]. Same memory layout,
+    so we build [B, HV, K] / [HV, K] and reshape for the baseline.
+
+We run a multi-step decode through BOTH kernels and assert:
+  * per-step outputs ``y`` match, AND
+  * the ground-truth full state matches the ReplaySSM checkpoint *after a flush*
+    (folding the ring buffer into S0). At a non-flush step the ReplaySSM
+    checkpoint legitimately lags (the update lives in the ring), so we only
+    compare state at flush boundaries / after a forced flush.
+
+L sweep:
+  * L=1 (flush every step): the ring is always empty and ``write_pos == L-1``
+    every step, so the reconstruction term is zero, the total decay is 1, and
+    the kernel reduces to the packed decode. For GDN the only fp difference is a
+    scalar ``alpha*(sum x)`` vs ``sum(alpha*x)`` reorder, so we assert ~bit-exact
+    (atol 2e-6). For KDA the analogous reorder is per-K (``alpha[k]`` folded into
+    q/k vs into S0), bounded by ~K*eps, so KDA uses the same (tighter-than-L>1)
+    tensor-core tolerances rather than 2e-6.
+  * L in {4,8,16}: algebraically equivalent but floating-point reordered (cumsum
+    of gates, replay-decay scaling, the d@k reconstruction dot on tensor cores
+    vs. the materialized recurrence). Tight fp tolerance, see ``_tols``.
+
+Runnable as ``pytest`` and as ``__main__``.
+"""
+
+import sys
+import unittest
+from pathlib import Path
+
+import torch
+
+from sglang.test.test_utils import CustomTestCase
+
+# Mirror sibling GDN unittests: register for CUDA/AMD CI. This is a kernel-math
+# unit test; it lives with the other linear-attention kernel correctness tests.
+# The registry calls MUST be module-level (the CI collector / check-registered-
+# tests hook parses them statically via AST and only scans top-level statements
+# -- a try/except wrapper hides them and fails the hook).
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
+register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-large-amd")
+
+
+def _build_inputs(B, H, HV, K, V, dtype, device, seed=0, is_kda=False):
+    """Random packed-decode inputs for one step.
+
+    Returns a packed ``mixed_qkv`` [B, 2*H*K + HV*V] plus the per-row gate
+    inputs. ``b`` (beta) is [B, HV] for both gate types. ``a`` is [B, HV] for
+    the GDN scalar gate and [B, HV, K] for the KDA per-K gate.
+    """
+    g = torch.Generator(device=device).manual_seed(seed)
+    qk_dim = 2 * H * K
+    v_dim = HV * V
+    mixed_qkv = torch.randn(B, qk_dim + v_dim, generator=g, device=device, dtype=dtype)
+    # Keep gate input `a` modest so softplus / decay stay in a stable range.
+    if is_kda:
+        a = (
+            torch.randn(B, HV, K, generator=g, device=device, dtype=dtype) * 0.5
+        ).contiguous()
+    else:
+        a = torch.randn(B, HV, generator=g, device=device, dtype=dtype) * 0.5
+    b = torch.randn(B, HV, generator=g, device=device, dtype=dtype)
+    return mixed_qkv, a, b
+
+
+def _tols(dtype):
+    # L>1 reconstructs the state via a tensor-core tl.dot (TF32 for fp32 inputs,
+    # bf16 tensor cores for bf16) -- the performance-critical path that keeps the
+    # L-deep reconstruction hidden under the reduced memory traffic. TF32 carries
+    # ~4e-4 error, bf16 ~1e-3, both benign end-to-end (ReplaySSM GSM8K parity);
+    # vLLM's own ReplaySSM test uses fp32 atol 1e-4. GDN L=1 has an empty buffer
+    # (dot == 0) and is asserted bit-exact separately.
+    if dtype == torch.float32:
+        return dict(atol=1e-4, rtol=1e-3)
+    return dict(atol=1e-3, rtol=1e-2)  # bf16
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
+class TestLinearReplaySSMDecode(CustomTestCase):
+    # Small config: HV=8 or 16 heads, K=V=128, batch B=4.
+    CONFIGS = (
+        dict(B=4, H=4, HV=8, K=128, V=128),
+        dict(B=4, H=8, HV=16, K=128, V=128),
+    )
+    NUM_STEPS = 40  # > 16 -> multiple flush cycles for L in {4,8,16}
+    L_SWEEP = (1, 4, 8, 16)
+
+    def _run_one(self, cfg, L, dtype, force_flush_steps=(), is_kda=False):
+        from sglang.srt.layers.attention.fla.fused_recurrent import (
+            fused_recurrent_gated_delta_rule_packed_decode,
+            fused_recurrent_kda_packed_decode,
+        )
+        from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+            fused_recurrent_linear_replayssm_decode,
+        )
+
+        device = "cuda"
+        B, H, HV, K, V = (cfg[x] for x in ("B", "H", "HV", "K", "V"))
+        scale = K**-0.5
+
+        # Shared static params across steps. A_log is per-head [HV] for both
+        # gate types; dt_bias is per-head [HV] (GDN) or per-K [HV, K] (KDA).
+        A_log = (torch.randn(HV, device=device, dtype=torch.float32) * 0.3).contiguous()
+        if is_kda:
+            dt_bias = (
+                torch.randn(HV, K, device=device, dtype=torch.float32) * 0.1
+            ).contiguous()
+        else:
+            dt_bias = (
+                torch.randn(HV, device=device, dtype=torch.float32) * 0.1
+            ).contiguous()
+
+        # Two independent state pools (same init), one per kernel. One slot/req.
+        num_slots = B
+        # Seed the state init so the marginal bf16 flush-state element is
+        # reproducible run-to-run (otherwise the pass/fail at the tolerance
+        # boundary flickers with the global RNG).
+        sg = torch.Generator(device=device).manual_seed(2024)
+        ref_state = torch.randn(
+            num_slots, HV, V, K, generator=sg, device=device, dtype=torch.float32
+        )
+        rep_state = ref_state.clone()
+        cache_indices = torch.arange(B, device=device, dtype=torch.int32)
+
+        # ReplaySSM ring tensors (caller-owned; this increment's caller == test).
+        # g_cache is per-step scalar [HV, L] for GDN, per-K [HV, L, K] for KDA.
+        d_cache = torch.zeros(num_slots, HV, L, V, device=device, dtype=dtype)
+        k_cache = torch.zeros(num_slots, H, L, K, device=device, dtype=dtype)
+        if is_kda:
+            g_cache = torch.zeros(
+                num_slots, HV, L, K, device=device, dtype=torch.float32
+            )
+        else:
+            g_cache = torch.zeros(num_slots, HV, L, device=device, dtype=torch.float32)
+        write_pos = torch.zeros(B, device=device, dtype=torch.int32)
+
+        tols = _tols(dtype)
+        flush_step_seen = False
+
+        for step in range(self.NUM_STEPS):
+            mixed_qkv, a, b = _build_inputs(
+                B, H, HV, K, V, dtype, device, seed=1000 + step, is_kda=is_kda
+            )
+
+            # ---- ground truth: packed decode, full-state write every step ----
+            ref_out = mixed_qkv.new_empty(B, 1, HV, V)
+            if is_kda:
+                # KDA baseline takes the per-K gate flattened: a=[B, HV*K],
+                # dt_bias=[HV*K] (same memory as our [B, HV, K] / [HV, K]).
+                fused_recurrent_kda_packed_decode(
+                    mixed_qkv=mixed_qkv,
+                    a=a.reshape(B, HV * K).contiguous(),
+                    b=b,
+                    A_log=A_log,
+                    dt_bias=dt_bias.reshape(HV * K).contiguous(),
+                    scale=scale,
+                    initial_state=ref_state,
+                    out=ref_out,
+                    ssm_state_indices=cache_indices,
+                    use_qk_l2norm_in_kernel=True,
+                )
+            else:
+                fused_recurrent_gated_delta_rule_packed_decode(
+                    mixed_qkv=mixed_qkv,
+                    a=a,
+                    b=b,
+                    A_log=A_log,
+                    dt_bias=dt_bias,
+                    scale=scale,
+                    initial_state=ref_state,
+                    out=ref_out,
+                    ssm_state_indices=cache_indices,
+                    use_qk_l2norm_in_kernel=True,
+                )
+
+            # ---- ReplaySSM buffered decode for this L ----
+            rep_out = mixed_qkv.new_empty(B, 1, HV, V)
+            # slice 2: force a flush at a radix track boundary even mid-buffer.
+            force_now = step in force_flush_steps
+            ff_tensor = (
+                torch.ones(B, device=device, dtype=torch.int32) if force_now else None
+            )
+            is_flush = bool((write_pos[0].item()) == L - 1) or force_now
+            flush_step_seen = flush_step_seen or is_flush
+            fused_recurrent_linear_replayssm_decode(
+                mixed_qkv=mixed_qkv,
+                a=a,
+                b=b,
+                A_log=A_log,
+                dt_bias=dt_bias,
+                scale=scale,
+                initial_state=rep_state,
+                d_cache=d_cache,
+                k_cache=k_cache,
+                g_cache=g_cache,
+                out=rep_out,
+                ssm_state_indices=cache_indices,
+                write_pos=write_pos,
+                force_flush=ff_tensor,
+                use_qk_l2norm_in_kernel=True,
+                is_kda=is_kda,
+                # nk=1 at L=1 keeps the K-reduction in a single tile (no
+                # K-tiling), matching the reference's single-pass reduction;
+                # nk=2 for L>1 exercises the K-tiled reconstruction path.
+                nk=1 if L == 1 else 2,
+            )
+
+            # Output must match every step.
+            if dtype == torch.float32 and L == 1 and not is_kda:
+                # GDN L=1: essentially bit-exact -- only an alpha*(sum) vs
+                # sum(alpha*) scalar-multiply reordering separates this from the
+                # reference. (KDA L=1 has the analogous per-K reorder, bounded by
+                # ~K*eps, so it uses the regular tensor-core tolerance below.)
+                torch.testing.assert_close(
+                    rep_out,
+                    ref_out,
+                    atol=2e-6,
+                    rtol=1e-5,
+                    msg=f"GDN L=1 step={step} out",
+                )
+            else:
+                torch.testing.assert_close(
+                    rep_out,
+                    ref_out,
+                    msg=f"is_kda={is_kda} L={L} step={step} out",
+                    **tols,
+                )
+
+            # On a flush step the ReplaySSM checkpoint is brought fully
+            # up-to-date and the ring logically clears -> compare states.
+            if is_flush:
+                if dtype == torch.float32 and L == 1 and not is_kda:
+                    torch.testing.assert_close(
+                        rep_state,
+                        ref_state,
+                        atol=2e-6,
+                        rtol=1e-5,
+                        msg=f"GDN L=1 step={step} state",
+                    )
+                else:
+                    # The flush state is reconstructed via the tensor-core dot,
+                    # so it carries tensor-core precision: ~tf32 (~4e-4, measured
+                    # max_abs ~7e-4) for fp32, ~bf16 (measured max_abs ~1.9e-3)
+                    # for bf16 -- looser than the per-step OUTPUT, which passes at
+                    # the tighter output tol because the tf32/bf16 error partly
+                    # cancels in the q-reduction. The per-step decay (alpha<1)
+                    # attenuates old reconstruction error so it does not
+                    # accumulate across flushes over a long decode; e2e accuracy
+                    # is unaffected (ReplaySSM GSM8K parity).
+                    state_tols = (
+                        dict(atol=1.5e-3, rtol=1e-3)
+                        if dtype == torch.float32
+                        else dict(atol=3e-3, rtol=1e-2)
+                    )
+                    torch.testing.assert_close(
+                        rep_state,
+                        ref_state,
+                        msg=f"is_kda={is_kda} L={L} step={step} state(flush)",
+                        **state_tols,
+                    )
+
+            # Advance the ring cursor (caller's responsibility in this phase):
+            # wrap to 0 after a flush (natural L-1 OR a forced track-boundary
+            # flush), otherwise increment.
+            if force_now:
+                write_pos = torch.zeros_like(write_pos)
+            else:
+                write_pos = torch.where(
+                    write_pos == L - 1,
+                    torch.zeros_like(write_pos),
+                    write_pos + 1,
+                )
+
+        # Explicit wrap/flush boundary must have been exercised for L>1; for
+        # L=1 every step is a flush.
+        self.assertTrue(flush_step_seen, f"flush boundary never hit for L={L}")
+
+    # ---- GDN (scalar gate) ----
+    def test_replayssm_matches_gdn_packed_decode_fp32(self):
+        for cfg in self.CONFIGS:
+            for L in self.L_SWEEP:
+                with self.subTest(cfg=cfg, L=L, dtype="fp32"):
+                    self._run_one(cfg, L, torch.float32)
+
+    def test_replayssm_matches_gdn_packed_decode_bf16(self):
+        for cfg in self.CONFIGS:
+            for L in self.L_SWEEP:
+                with self.subTest(cfg=cfg, L=L, dtype="bf16"):
+                    self._run_one(cfg, L, torch.bfloat16)
+
+    # ---- KDA (per-K gate) ----
+    def test_replayssm_matches_kda_packed_decode_fp32(self):
+        for cfg in self.CONFIGS:
+            for L in self.L_SWEEP:
+                with self.subTest(cfg=cfg, L=L, dtype="fp32"):
+                    self._run_one(cfg, L, torch.float32, is_kda=True)
+
+    def test_replayssm_matches_kda_packed_decode_bf16(self):
+        for cfg in self.CONFIGS:
+            for L in self.L_SWEEP:
+                with self.subTest(cfg=cfg, L=L, dtype="bf16"):
+                    self._run_one(cfg, L, torch.bfloat16, is_kda=True)
+
+    # ---- forced (radix track-boundary) flush, both gate types ----
+    def test_force_flush_matches_reference(self):
+        """A radix track-boundary forces a flush mid-buffer (write_pos != L-1).
+        The kernel must fold the partial ring + current token into the
+        checkpoint so an external snapshot reads an up-to-date state; output
+        every step and the forced-flush checkpoint must still match the running
+        recurrence -- for both the GDN scalar gate and the KDA per-K gate."""
+        cfg = dict(B=4, H=8, HV=16, K=128, V=128)
+        for is_kda in (False, True):
+            for dtype in (torch.float32, torch.bfloat16):
+                with self.subTest(is_kda=is_kda, dtype=str(dtype)):
+                    self._run_one(
+                        cfg,
+                        L=8,
+                        dtype=dtype,
+                        force_flush_steps=(3, 5, 11, 20),
+                        is_kda=is_kda,
+                    )
+
+    def test_flush_boundary_state_exact_l1(self):
+        """At L=1 (GDN) every step is a flush; the kernel is algebraically the
+        packed decode. Each step reproduces the reference output AND the full
+        state to atol 2e-6 / rtol 1e-5 (essentially bit-exact) in fp32. This also
+        exercises the ``write_pos == L-1`` flush boundary on every step."""
+        self._run_one(dict(B=4, H=4, HV=8, K=128, V=128), L=1, dtype=torch.float32)
+
+
+if __name__ == "__main__":
+    unittest.main()