diff --git a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py index fa6150eafd25..e4342cbd9882 100644 --- a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py +++ b/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 diff --git a/test/registered/attention/unittests/gdn/test_gdn_replayssm_compact_commit.py b/test/registered/attention/unittests/gdn/test_gdn_replayssm_compact_commit.py new file mode 100644 index 000000000000..45f2e7e75478 --- /dev/null +++ b/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"]))