[Fix] Mask unused ReplaySSM history before checkpoint arithmetic

python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py

@@ -673,6 +673,8 @@ def gdn_replayssm_compact_commit_kernel(
     base = tl.load(cache_base + replay_idx).to(tl.int32)
     o_t = tl.arange(0, MAX_CACHE_LEN)
     phys = (base + o_t) & (MAX_CACHE_LEN - 1)
+    # Zero decay cannot suppress NaN/Inf left in unused ring slots.
+    history_mask = o_t < n_history
     o_k = tl.arange(0, K)
     o_v = i_v * BV + tl.arange(0, BV)
     mask_v = o_v < V
@@ -689,21 +691,25 @@ def gdn_replayssm_compact_commit_kernel(
         k_cache
         + replay_idx * stride_k_slot
         + (i_h * MAX_CACHE_LEN + phys[:, None]) * K
-        + o_k[None, :]
+        + o_k[None, :],
+        mask=history_mask[:, None],
+        other=0.0,
     )
     if HAS_RESIDUAL:
         key_residual = tl.load(
             k_residual_cache
             + replay_idx * stride_k_residual_slot
             + (i_h * MAX_CACHE_LEN + phys[:, None]) * K
-            + o_k[None, :]
+            + o_k[None, :],
+            mask=history_mask[:, None],
+            other=0.0,
         )
     updates = tl.load(
         d_cache
         + replay_idx * stride_d_slot
         + (i_hv * MAX_CACHE_LEN + phys[:, None]) * V
         + o_v[None, :],
-        mask=mask_v[None, :],
+        mask=history_mask[:, None] & mask_v[None, :],
         other=0.0,
     ).to(tl.float32)
     if HAS_RESIDUAL:
@@ -712,19 +718,19 @@ def gdn_replayssm_compact_commit_kernel(
             + replay_idx * stride_d_residual_slot
             + (i_hv * MAX_CACHE_LEN + phys[:, None]) * V
             + o_v[None, :],
-            mask=mask_v[None, :],
+            mask=history_mask[:, None] & mask_v[None, :],
             other=0.0,
         ).to(tl.float32)
     gates = tl.load(
-        g_cache + replay_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys
+        g_cache + replay_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys,
+        mask=history_mask,
+        other=0.0,
     ).to(tl.float32)
 
     if fold_active:
-        active_mask = o_t < n_history
-        active_g = tl.where(active_mask, gates, 0.0)
-        active_prefix = tl.cumsum(active_g, axis=0)
-        active_total = tl.sum(active_g, axis=0)
-        active_decay = tl.where(active_mask, tl.exp(active_total - active_prefix), 0.0)
+        active_prefix = tl.cumsum(gates, axis=0)
+        active_total = tl.sum(gates, axis=0)
+        active_decay = tl.where(history_mask, tl.exp(active_total - active_prefix), 0.0)
         active_updates = updates * active_decay[:, None]
         active_hi = active_updates.to(k_cache.dtype.element_ty)
         active_lo = (active_updates - active_hi.to(tl.float32)).to(
@@ -748,16 +754,19 @@ def gdn_replayssm_compact_commit_kernel(
             track_prefix = tl.cumsum(track_g, axis=0)
             track_total = tl.sum(track_g, axis=0)
             track_decay = tl.where(track_mask, tl.exp(track_total - track_prefix), 0.0)
-            track_updates = updates * track_decay[:, None]
+            track_keys = tl.where(track_mask[:, None], keys, 0.0)
+            track_updates = tl.where(track_mask[:, None], updates, 0.0)
+            track_updates *= track_decay[:, None]
             track_hi = track_updates.to(k_cache.dtype.element_ty)
             track_lo = (track_updates - track_hi.to(tl.float32)).to(
                 k_cache.dtype.element_ty
             )
-            track_delta = tl.dot(tl.trans(keys), track_hi)
-            track_delta += tl.dot(tl.trans(keys), track_lo)
+            track_delta = tl.dot(tl.trans(track_keys), track_hi)
+            track_delta += tl.dot(tl.trans(track_keys), track_lo)
             if HAS_RESIDUAL:
-                track_delta += tl.dot(tl.trans(key_residual), track_hi)
-                track_delta += tl.dot(tl.trans(key_residual), track_lo)
+                track_key_residual = tl.where(track_mask[:, None], key_residual, 0.0)
+                track_delta += tl.dot(tl.trans(track_key_residual), track_hi)
+                track_delta += tl.dot(tl.trans(track_key_residual), track_lo)
             tl.store(
                 h0
                 + track_idx * stride_state_slot

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

@@ -0,0 +1,155 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_decode import (
+    commit_gdn_replayssm_circular,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large")
+
+DEVICE = "cuda"
+CACHE_LEN = 16
+HISTORY_LEN = 4
+TRACK_LEN = 3
+LAYERS, SLOTS, H, HV, K, V = 2, 6, 1, 2, 64, 80
+FIELDS = ("d_cache", "k_cache", "g_cache", "d_residual_cache", "k_residual_cache")
+
+
+def _run_commit(
+    *,
+    base: int,
+    mode: str,
+    residual: bool,
+    field: str | None = None,
+    poison: float = 0.0,
+    poison_start: int = HISTORY_LEN,
+    history_len: int = HISTORY_LEN,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    gen = torch.Generator(device=DEVICE).manual_seed(42)
+
+    def rand(*shape: int, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        return torch.randn(*shape, generator=gen, device=DEVICE, dtype=dtype) * 0.125
+
+    state = rand(LAYERS, SLOTS, HV, V, K, dtype=torch.float32)
+    initial = state.clone()
+    caches = {
+        "d_cache": rand(LAYERS, SLOTS, HV, CACHE_LEN, V),
+        "k_cache": rand(LAYERS, SLOTS, H, CACHE_LEN, K),
+        "g_cache": torch.full((LAYERS, SLOTS, HV, CACHE_LEN), -0.03125, device=DEVICE),
+        "d_residual_cache": rand(LAYERS, SLOTS, HV, CACHE_LEN, V) * 0.001,
+        "k_residual_cache": rand(LAYERS, SLOTS, H, CACHE_LEN, K) * 0.001,
+    }
+    physical = (base + torch.arange(CACHE_LEN, device=DEVICE)) % CACHE_LEN
+    if field is not None:
+        caches[field][:, 2, :, physical[poison_start:]] = poison
+
+    def ints(values: list[int]) -> torch.Tensor:
+        return torch.tensor(values, device=DEVICE, dtype=torch.int32)
+
+    state_indices = ints([1, 2, 0])
+    replay_indices = ints([2, 1, 0])
+    write_pos = ints([history_len, history_len, history_len])
+    cache_base = ints([base, base, base])
+    is_flush = ints([1, 0, int(mode != "track")])
+    accepted = ints([2, 2, 2])
+    track_indices = ints([3, 0, 5]) if mode != "active" else None
+    track_steps = ints([0, -1, 0]) if mode != "active" else None
+
+    commit_gdn_replayssm_circular(
+        checkpoint_state=state,
+        d_cache=caches["d_cache"],
+        k_cache=caches["k_cache"],
+        g_cache=caches["g_cache"],
+        d_residual_cache=caches["d_residual_cache"] if residual else None,
+        k_residual_cache=caches["k_residual_cache"] if residual else None,
+        state_batch_indices=state_indices,
+        replay_indices=replay_indices,
+        write_pos=write_pos,
+        cache_base=cache_base,
+        is_flush=is_flush,
+        accept_lens=accepted,
+        mamba_track_indices=track_indices,
+        mamba_steps_to_track=track_steps,
+        null_block_id=0,
+    )
+
+    changed = set()
+    if mode != "track":
+        changed.add(1)
+    if mode != "active":
+        changed.add(3)
+    for slot in set(range(SLOTS)) - changed:
+        torch.testing.assert_close(state[:, slot], initial[:, slot], rtol=0, atol=0)
+    torch.testing.assert_close(
+        write_pos,
+        ints([history_len, history_len, history_len if mode == "track" else 0]),
+        rtol=0,
+        atol=0,
+    )
+    torch.testing.assert_close(
+        cache_base,
+        ints(
+            [base, base, base if mode == "track" else (base + history_len) % CACHE_LEN]
+        ),
+        rtol=0,
+        atol=0,
+    )
+    torch.testing.assert_close(is_flush, ints([1, 0, 0]), rtol=0, atol=0)
+    return state, initial
+
+
+@pytest.mark.parametrize(
+    "field,residual",
+    [(field, True) for field in FIELDS] + [(field, False) for field in FIELDS[:3]],
+)
+@pytest.mark.parametrize("poison", [float("nan"), float("inf"), -float("inf")])
+@pytest.mark.parametrize(
+    "mode,base", [("active", 0), ("track", 13), ("both", 0), ("both", 13)]
+)
+def test_unused_history_cannot_contaminate_checkpoint(
+    field: str, residual: bool, poison: float, mode: str, base: int
+) -> None:
+    clean, _ = _run_commit(base=base, mode=mode, residual=residual)
+    reused, _ = _run_commit(
+        base=base, mode=mode, residual=residual, field=field, poison=poison
+    )
+    assert torch.isfinite(reused).all()
+    torch.testing.assert_close(reused, clean, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("field", FIELDS)
+@pytest.mark.parametrize("base", [0, 13])
+def test_tracked_prefix_excludes_later_valid_nan(field: str, base: int) -> None:
+    clean, _ = _run_commit(base=base, mode="both", residual=True)
+    poisoned, _ = _run_commit(
+        base=base,
+        mode="both",
+        residual=True,
+        field=field,
+        poison=float("nan"),
+        poison_start=TRACK_LEN,
+    )
+    assert not torch.isfinite(poisoned[:, 1]).all()
+    assert torch.isfinite(poisoned[:, 3]).all()
+    torch.testing.assert_close(poisoned[:, 3], clean[:, 3], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("field", FIELDS)
+def test_empty_history_preserves_checkpoint(field: str) -> None:
+    result, initial = _run_commit(
+        base=13,
+        mode="active",
+        residual=True,
+        field=field,
+        poison=float("nan"),
+        poison_start=0,
+        history_len=0,
+    )
+    torch.testing.assert_close(result, initial, rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))