diff --git a/test/registered/kernels/test_kda_mtp_cutedsl_replayssm_ring.py b/test/registered/kernels/test_kda_mtp_cutedsl_replayssm_ring.py new file mode 100644 index 000000000..f83aa6a9a --- /dev/null +++ b/test/registered/kernels/test_kda_mtp_cutedsl_replayssm_ring.py @@ -0,0 +1,194 @@ +"""Parity: CuTe MTP verify kernel's ReplaySSM ring == its own state snapshots. + +`fused_kda_decode_mtp_dspark(replayssm_*=...)` switches the CuTe DSpARK verify +kernel to CACHE_RING mode: per draft step it stores post-conv pre-l2norm k, +post-conv v, the log-decay gate gk, and sigmoid(beta) into the per-slot rings +(and skips the per-step intermediate_ssm snapshots). Each case runs the kernel +twice on identical inputs — baseline arm producing snapshots, ring arm +producing rings — folds the ring with `commit_kda_replayssm_spec`, and checks +the folded checkpoint against the baseline arm's last-step snapshot. A wrong +head/step/slot offset in the fused ring store shows up as a mismatch; the +output tensor must be untouched by the mode (bitwise equal). + +Tolerance is bf16-bound: production rings store rawk/rawv in conv dtype +(bf16), so the fold re-quantizes k/v while the baseline snapshot keeps them +in fp32 registers. +""" + +import pytest +import torch + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +# SM100-only, and hard: libNVVM refuses the generated device IR when the CuTe +# DSL kernel is compiled for sm_90a, so this is a build failure rather than a +# numeric one. The SM100 pool has no single-GPU runner_config, hence 4-gpu-b200 +# for a one-GPU test. +if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + pytest.skip( + "KDA CuTe MTP ReplaySSM ring parity needs SM100 (CuTe DSL kernel).", + allow_module_level=True, + ) + +from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( # noqa: E402 + commit_kda_replayssm_spec, +) +from sglang.kernels.ops.kimi_k3.kda_decode_mtp import ( # noqa: E402 + fused_kda_decode_mtp_dspark, +) + +DEV = "cuda" +K = 128 +W = 4 # KERNEL_WIDTH + + +def _run(arm, *, N, H, num_spec, seed, onorm=False, pad_last=False): + torch.manual_seed(seed) + T = N * (1 + num_spec) + num_slots = N + 2 + L = 16 + assert L >= 1 + num_spec + + x_q = torch.randn(1, T, H, K, device=DEV, dtype=torch.bfloat16) + x_k = torch.randn(1, T, H, K, device=DEV, dtype=torch.bfloat16) + x_v = torch.randn(1, T, H, K, device=DEV, dtype=torch.bfloat16) + g = torch.randn(1, T, H, K, device=DEV, dtype=torch.bfloat16) + beta = torch.randn(1, T, H, device=DEV, dtype=torch.bfloat16) + w_q = torch.randn(H * K, W, device=DEV, dtype=torch.float32) * 0.1 + w_k = torch.randn(H * K, W, device=DEV, dtype=torch.float32) * 0.1 + w_v = torch.randn(H * K, W, device=DEV, dtype=torch.float32) * 0.1 + cs_q = torch.randn(num_slots, H * K, W - 1, device=DEV, dtype=torch.bfloat16) + cs_k = torch.randn(num_slots, H * K, W - 1, device=DEV, dtype=torch.bfloat16) + cs_v = torch.randn(num_slots, H * K, W - 1, device=DEV, dtype=torch.bfloat16) + A_log = torch.randn(H, device=DEV, dtype=torch.float32) + dt_bias = torch.randn(H * K, device=DEV, dtype=torch.float32) + h0 = torch.randn(num_slots, H, K, K, device=DEV, dtype=torch.float32) + + slots = torch.arange(1, N + 1, device=DEV, dtype=torch.int32) + if pad_last: + slots[-1] = -1 + scratch = torch.arange(N, device=DEV, dtype=torch.int32) + cu_seqlens = torch.arange(0, T + 1, 1 + num_spec, device=DEV, dtype=torch.int32) + + ic_q = torch.zeros(N, 1 + num_spec, H * K, W - 1, device=DEV, dtype=torch.bfloat16) + ic_k = torch.zeros_like(ic_q) + ic_v = torch.zeros_like(ic_q) + + kwargs = dict( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=h0, + intermediate_state_indices=scratch, + intermediate_conv_q=ic_q, + intermediate_conv_k=ic_k, + intermediate_conv_v=ic_v, + ssm_state_indices=slots, + cu_seqlens=cu_seqlens, + lower_bound=-5.0, + ) + norm = {} + if onorm: + norm = dict( + gate=torch.randn(1, T, H, K, device=DEV, dtype=torch.bfloat16), + weight=torch.randn(K, device=DEV, dtype=torch.float32), + eps=1e-6, + ) + kwargs.update( + onorm_gate=norm["gate"], + onorm_weight=norm["weight"], + onorm_eps=norm["eps"], + ) + if arm == "baseline": + inter = torch.zeros(N, 1 + num_spec, H, K, K, device=DEV, dtype=torch.float32) + out = fused_kda_decode_mtp_dspark(intermediate_ssm=inter, **kwargs) + return out, dict(inter=inter, slots=slots, scratch=scratch, **norm) + rawv = torch.zeros(num_slots, H, L, K, device=DEV, dtype=torch.bfloat16) + rawk = torch.zeros_like(rawv) + gring = torch.zeros(num_slots, H, L, K, device=DEV, dtype=torch.float32) + betar = torch.zeros(num_slots, H, L, device=DEV, dtype=torch.float32) + out = fused_kda_decode_mtp_dspark( + intermediate_ssm=None, + replayssm_rawv=rawv, + replayssm_rawk=rawk, + replayssm_g=gring, + replayssm_beta=betar, + **kwargs, + ) + return out, dict(rawv=rawv, rawk=rawk, gring=gring, betar=betar, h0=h0, slots=slots) + + +@pytest.mark.parametrize( + "N,H,num_spec", [(4, 2, 4), (1, 12, 5), (16, 2, 8)], ids=["small", "k3ish", "wide"] +) +def test_cutedsl_ring_fold_parity(N, H, num_spec): + seed = 0 + out_base, base = _run("baseline", N=N, H=H, num_spec=num_spec, seed=seed) + out_ring, ring = _run("ring", N=N, H=H, num_spec=num_spec, seed=seed) + + torch.testing.assert_close(out_ring, out_base, rtol=0, atol=0) + + T_req = 1 + num_spec + acc = torch.full((N,), T_req, device=DEV, dtype=torch.int32) + ckpt = ring["h0"].clone() + commit_kda_replayssm_spec( + ckpt, + ring["rawv"], + ring["rawk"], + ring["gring"], + ring["betar"], + ring["slots"], + acc, + max_cache_len=ring["rawv"].shape[2], + num_k_heads=H, + use_qk_l2norm_in_kernel=True, + null_block_id=-1, + ) + for j in range(N): + base_state = base["inter"][base["scratch"][j], T_req - 1] + fold = ckpt[ring["slots"][j]] + rel = ( + (fold - base_state).abs().max() / base_state.abs().max().clamp_min(1e-6) + ).item() + assert rel < 2e-2, f"req={j}: rel={rel:.3e}" + + +@pytest.mark.parametrize("N", [4, 32], ids=["small-grid", "large-grid"]) +def test_cutedsl_fused_output_norm(N): + H, num_spec, seed = 2, 2, 7 + raw, _ = _run("baseline", N=N, H=H, num_spec=num_spec, seed=seed) + fused, norm = _run("baseline", N=N, H=H, num_spec=num_spec, seed=seed, onorm=True) + + ref = raw.float() + ref = ref * torch.rsqrt(ref.square().mean(dim=-1, keepdim=True) + norm["eps"]) + ref = ref * norm["weight"] * torch.sigmoid(norm["gate"].float()) + torch.testing.assert_close(fused.float(), ref, rtol=2e-2, atol=3e-2) + + +@pytest.mark.parametrize("N", [4, 32], ids=["small-grid", "large-grid"]) +def test_cutedsl_cuda_graph_padding_slot_is_safe(N): + _, ring = _run("ring", N=N, H=2, num_spec=2, seed=11, pad_last=True) + torch.cuda.synchronize() + + # The last logical request is graph padding. Its original physical slot N + # is now unused and must remain untouched by all ReplaySSM ring stores. + for name in ("rawv", "rawk", "gring", "betar"): + assert torch.count_nonzero(ring[name][N]).item() == 0 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/test_kda_replayssm_fold.py b/test/registered/kernels/test_kda_replayssm_fold.py new file mode 100644 index 000000000..249b40076 --- /dev/null +++ b/test/registered/kernels/test_kda_replayssm_fold.py @@ -0,0 +1,147 @@ +"""Parity: KDA ReplaySSM fold-commit vs the recurrent verify kernel. + +The fold kernel (`kda_replayssm_exact_fold_kernel`) replays a request's accepted +draft window from a checkpoint to reconstruct the committed SSM state, replacing +the per-step `intermediate_ssm` snapshots. This test pins the fold's committed +state to the recurrent verify kernel's state (derived-property parity). + +Crucially it drives BOTH sides from raw (a, b, A_log, dt_bias): + - baseline: the verify kernel (`fused_sigmoid_gating_delta_rule_update`, is_kda) + forms the gate INTERNALLY and caches per-step states. + - fold: the backend forms gk/beta in torch, writes them to the ring, and the + fold kernel replays them. +So a gate-formula mismatch is caught here. This guards a real bug: the fold +originally formed the gate with plain `softplus` while K3's checkpoint uses the +safe gate (`lower_bound * sigmoid(exp(A_log) * x)`, gate_lower_bound=-5.0), +which silently committed the wrong state (gsm8k 0.955 -> 0.947). The earlier +same-gk fold test could not catch it because it fed both sides the same gk. +Running the safe-gate case on the pre-fix code turns it red. +""" + +import unittest + +import torch + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +if not torch.cuda.is_available(): + import pytest + + pytest.skip( + "KDA ReplaySSM fold parity needs CUDA (triton kernels).", + allow_module_level=True, + ) + +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( # noqa: E402 + fused_sigmoid_gating_delta_rule_update, +) +from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( # noqa: E402 + commit_kda_replayssm_spec, +) + + +class TestKDAReplaySSMFoldParity(CustomTestCase): + B, T, HV, K, V = 3, 6, 4, 64, 64 + H = HV # single k-head group + L = 16 + + def _parity(self, lower_bound): + dev = "cuda" + B, T, HV, K, V, H, L = self.B, self.T, self.HV, self.K, self.V, self.H, self.L + scale = K**-0.5 + torch.manual_seed(0) + + q = torch.randn(B, T, H, K, device=dev, dtype=torch.float32) + k = torch.randn(B, T, H, K, device=dev, dtype=torch.float32) + v = torch.randn(B, T, HV, V, device=dev, dtype=torch.float32) + a = torch.randn(B, T, HV, K, device=dev, dtype=torch.float32) + b = torch.randn(B, T, HV, device=dev, dtype=torch.float32) + A_log = torch.randn(HV, device=dev, dtype=torch.float32) + dt_bias = torch.randn(HV, K, device=dev, dtype=torch.float32) + h0 = torch.randn(B, HV, V, K, device=dev, dtype=torch.float32) + + slots = torch.arange(1, B + 1, device=dev, dtype=torch.int32) + num_slots = B + 1 + + # baseline: verify kernel forms the gate internally, caches per-step state + h0_src = torch.zeros(num_slots, HV, V, K, device=dev, dtype=torch.float32) + for j in range(B): + h0_src[slots[j]] = h0[j] + inter = torch.zeros(num_slots, T, HV, V, K, device=dev, dtype=torch.float32) + fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=h0_src, + initial_state_indices=slots, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + lower_bound=lower_bound, + disable_state_update=True, + intermediate_states_buffer=inter, + intermediate_state_indices=slots, + cache_steps=T, + ) + accept = T # commit the full window + base = torch.stack([inter[slots[j], accept - 1] for j in range(B)], 0) + + # fold: torch forms gk/beta (matching the kernel's two branches), replays + x = a + dt_bias.view(1, 1, HV, K) + exp_a_log = torch.exp(A_log).view(1, 1, HV, 1) + if lower_bound is not None: + gk = lower_bound * torch.sigmoid(exp_a_log * x) + else: + gk = -exp_a_log * torch.nn.functional.softplus(x) + beta = torch.sigmoid(b) + + rawv = torch.zeros(num_slots, HV, L, V, device=dev, dtype=torch.float32) + rawk = torch.zeros(num_slots, H, L, K, device=dev, dtype=torch.float32) + gkr = torch.zeros(num_slots, HV, L, K, device=dev, dtype=torch.float32) + betar = torch.zeros(num_slots, HV, L, device=dev, dtype=torch.float32) + ckpt = torch.zeros(num_slots, HV, V, K, device=dev, dtype=torch.float32) + for j in range(B): + s = slots[j].item() + rawv[s, :, :T] = v[j].transpose(0, 1) + rawk[s, :, :T] = k[j].transpose(0, 1) + gkr[s, :, :T] = gk[j].transpose(0, 1) + betar[s, :, :T] = beta[j].transpose(0, 1) + ckpt[s] = h0[j] + acc = torch.full((B,), accept, device=dev, dtype=torch.int32) + commit_kda_replayssm_spec( + ckpt, + rawv, + rawk, + gkr, + betar, + slots, + acc, + max_cache_len=L, + num_k_heads=H, + use_qk_l2norm_in_kernel=True, + ) + fold = torch.stack([ckpt[slots[j].item()] for j in range(B)], 0) + + rel = ((fold - base).abs().max() / base.abs().max().clamp_min(1e-6)).item() + self.assertLess(rel, 1e-3, f"fold vs verify parity failed: rel={rel:.3e}") + + def test_safe_gate(self): + # K3's gate: g = lower_bound * sigmoid(exp(A_log) * (a + dt_bias)). + self._parity(lower_bound=-5.0) + + def test_softplus_gate(self): + # Plain branch: g = -exp(A_log) * softplus(a + dt_bias). + self._parity(lower_bound=None) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/test_kda_replayssm_fold_batched.py b/test/registered/kernels/test_kda_replayssm_fold_batched.py new file mode 100644 index 000000000..938c7b1d4 --- /dev/null +++ b/test/registered/kernels/test_kda_replayssm_fold_batched.py @@ -0,0 +1,140 @@ +"""Parity: layer-batched fold == per-layer loop of commit_kda_replayssm_spec. + +The commit-side ReplaySSM fold used to launch commit_kda_replayssm_spec once per +KDA layer (a Python loop -> ~69 tiny eager launches at bs=1, dispatch-bound). +`commit_kda_replayssm_spec_all_layers` packs the layer into the head grid axis so +one launch folds every layer. This must be BIT-IDENTICAL to the loop: each +(layer, head, v-tile) block runs the same per-slot recurrent replay; batching only +fans out the grid. Any layer-offset / stride bug shows up as a per-layer mismatch. + +Shapes cover GQA (H != HV), non-pow2 K/V, single accept step, padding (-1) slots, +varied accept_lens (incl. 0 = skip), track on/off, and single-layer (i_layer==0 +must match the per-layer entry). GPU-only. +""" + +import pytest +import torch + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +if not torch.cuda.is_available(): + pytest.skip( + "KDA ReplaySSM batched-fold parity needs CUDA (triton).", + allow_module_level=True, + ) + +from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( # noqa: E402 + commit_kda_replayssm_spec, + commit_kda_replayssm_spec_all_layers, +) + +DEV = "cuda" + +# (num_layers, num_slots, HV, H, K, V, L, has_track) +SHAPES = [ + (3, 6, 4, 4, 64, 64, 16, False), # square heads, no track + (2, 5, 8, 2, 128, 128, 16, True), # GQA 4:1, track on + (4, 7, 6, 3, 96, 80, 16, False), # non-pow2 K/V, GQA 2:1 + (2, 4, 4, 4, 64, 64, 16, True), # track on + (1, 4, 4, 4, 64, 64, 8, False), # single layer -> i_layer==0 path +] +SHAPE_IDS = ["square", "gqa4-track", "nonpow2", "track", "single-layer"] + + +def _rand_rings(num_layers, num_slots, HV, H, K, V, L): + torch.manual_seed(0) + return dict( + temporal=torch.randn( + num_layers, num_slots, HV, V, K, device=DEV, dtype=torch.float32 + ), + rawv=torch.randn( + num_layers, num_slots, HV, L, V, device=DEV, dtype=torch.bfloat16 + ), + rawk=torch.randn( + num_layers, num_slots, H, L, K, device=DEV, dtype=torch.bfloat16 + ), + gk=(-5.0) + * torch.sigmoid( + torch.randn( + num_layers, num_slots, HV, L, K, device=DEV, dtype=torch.float32 + ) + ), + beta=torch.sigmoid( + torch.randn(num_layers, num_slots, HV, L, device=DEV, dtype=torch.float32) + ), + ) + + +@pytest.mark.parametrize( + "num_layers,num_slots,HV,H,K,V,L,has_track", SHAPES, ids=SHAPE_IDS +) +def test_fold_batched_matches_per_layer( + num_layers, num_slots, HV, H, K, V, L, has_track +): + r = _rand_rings(num_layers, num_slots, HV, H, K, V, L) + B = min(3, num_slots - 1) + # slots 1..B; one row is a -1 padding slot (must be skipped identically). + slots = torch.arange(1, B + 1, device=DEV, dtype=torch.int32) + if B > 1: + slots[-1] = -1 + # accept_lens: mix 0 (skip), mid, full L. + accept = torch.tensor( + [[0, L // 2, L][i % 3] for i in range(B)], device=DEV, dtype=torch.int32 + ) + if has_track: + track_idx = torch.arange( + num_slots - B, num_slots, device=DEV, dtype=torch.int32 + ) + track_step = torch.tensor( + [[-1, L // 2 - 1, 1][i % 3] for i in range(B)], + device=DEV, + dtype=torch.int32, + ) + else: + track_idx = track_step = None + + # reference: per-layer loop + ref = r["temporal"].clone() + for li in range(num_layers): + commit_kda_replayssm_spec( + checkpoint_state=ref[li], + rawv_cache=r["rawv"][li], + rawk_cache=r["rawk"][li], + gk_cache=r["gk"][li], + beta_cache=r["beta"][li], + ssm_state_indices=slots, + accept_lens=accept, + max_cache_len=L, + num_k_heads=H, + mamba_track_indices=track_idx, + mamba_steps_to_track=track_step, + null_block_id=-1, + ) + # batched: one launch + bat = r["temporal"].clone() + commit_kda_replayssm_spec_all_layers( + checkpoint_state=bat, + rawv_cache=r["rawv"], + rawk_cache=r["rawk"], + gk_cache=r["gk"], + beta_cache=r["beta"], + ssm_state_indices=slots, + accept_lens=accept, + max_cache_len=L, + num_k_heads=H, + mamba_track_indices=track_idx, + mamba_steps_to_track=track_step, + null_block_id=-1, + ) + assert torch.equal(ref, bat), ( + f"batched fold != per-layer loop; max abs diff " + f"{(ref - bat).abs().max().item():.3e}" + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/test_kda_replayssm_ring_fused.py b/test/registered/kernels/test_kda_replayssm_ring_fused.py new file mode 100644 index 000000000..afbba8dbb --- /dev/null +++ b/test/registered/kernels/test_kda_replayssm_ring_fused.py @@ -0,0 +1,145 @@ +"""Parity: CACHE_RING verify kernel's ring == the state the verify kernel commits. + +This fuses the ReplaySSM ring-write into the recurrent verify kernel +(`fused_sigmoid_gating_delta_rule_update`, CACHE_RING=True): every draft step it +stores the pre-norm k / raw v / in-kernel gate / beta into the per-slot ring, in +place of the eager torch ring-write in kda_backend. Each case drives the fused +kernel end-to-end: run verify with CACHE_RING=True to fill the ring, fold the ring +back, and check the folded checkpoint matches the verify kernel's own per-step +state (intermediate_states_buffer). The ring's gate comes from the kernel's +tl.sigmoid/tl.exp (same as the state update), so the fold is bit-close on shape. + +Shapes guard the fused store's head-index / tile / step offsets: GQA (H != HV), +non-pow2 K/V, T 1: + slots[-1] = -1 + num_slots = bs + 1 + + h0_src = torch.zeros(num_slots, HV, V, K, device=DEV, dtype=torch.float32) + for j in range(bs): + if slots[j] >= 0: + h0_src[slots[j]] = h0[j] + inter = torch.zeros(num_slots, T, HV, V, K, device=DEV, dtype=torch.float32) + + # ring buffers filled by the fused kernel (CACHE_RING=True). + rawv = torch.zeros(num_slots, HV, L, V, device=DEV, dtype=torch.float32) + rawk = torch.zeros(num_slots, H, L, K, device=DEV, dtype=torch.float32) + gring = torch.zeros(num_slots, HV, L, K, device=DEV, dtype=torch.float32) + betar = torch.zeros(num_slots, HV, L, device=DEV, dtype=torch.float32) + + fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=h0_src, + initial_state_indices=slots, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + lower_bound=lower_bound, + disable_state_update=True, + intermediate_states_buffer=inter, + intermediate_state_indices=slots, + cache_steps=T, + # fused ring-write (kwargs below are added by the fusion). + cache_ring=True, + replayssm_rawv=rawv, + replayssm_rawk=rawk, + replayssm_g=gring, + replayssm_beta=betar, + ) + + acc = torch.full((bs,), T, device=DEV, dtype=torch.int32) + ckpt = h0_src.clone() + commit_kda_replayssm_spec( + ckpt, + rawv, + rawk, + gring, + betar, + slots, + acc, + max_cache_len=L, + num_k_heads=H, + use_qk_l2norm_in_kernel=True, + null_block_id=-1, + ) + + for j in range(bs): + if slots[j] < 0: + continue # padding row: ring not written, nothing to check + base = inter[slots[j], T - 1] + fold = ckpt[slots[j]] + rel = ((fold - base).abs().max() / base.abs().max().clamp_min(1e-6)).item() + assert rel < 1e-3, f"row={j}: rel={rel:.3e}" + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/test_kda_replayssm_ring_ragged.py b/test/registered/kernels/test_kda_replayssm_ring_ragged.py new file mode 100644 index 000000000..b428b263e --- /dev/null +++ b/test/registered/kernels/test_kda_replayssm_ring_ragged.py @@ -0,0 +1,150 @@ +"""Parity: CACHE_RING ring-write under ragged (varlen) verify layouts. + +Packed varlen verify (per-row verify_lens <= gamma): the fold of the first +acc entries of each row's ring must match the kernel's own per-step state, +as in the dense parity test. Covers partial commit (the compact commit shape) +and padding (-1) slots. GPU-only. +""" + +import pytest +import torch + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +if not torch.cuda.is_available(): + pytest.skip( + "KDA ReplaySSM ragged ring-write parity needs CUDA (triton).", + allow_module_level=True, + ) + +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( # noqa: E402 + fused_sigmoid_gating_delta_rule_update, +) +from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( # noqa: E402 + commit_kda_replayssm_spec, +) + +DEV = "cuda" + +# (bs, gamma, HV, H, K, V) +SHAPES = [ + (16, 8, 4, 4, 64, 64), # baseline square heads + (8, 8, 32, 32, 128, 128), # K3-like TP8 shape + (8, 4, 8, 2, 128, 128), # GQA: 4 v-heads per k-head +] +SHAPE_IDS = ["square", "k3-tp8", "gqa4"] + + +def _run_case(bs, gamma, HV, H, K, V, lens, acc, L, pad_last=False): + scale = K**-0.5 + total = int(lens.sum()) + cu = torch.zeros(bs + 1, device=DEV, dtype=torch.int32) + cu[1:] = torch.cumsum(lens, dim=0) + + q = torch.randn(1, total, H, K, device=DEV, dtype=torch.float32) + k = torch.randn(1, total, H, K, device=DEV, dtype=torch.float32) + v = torch.randn(1, total, HV, V, device=DEV, dtype=torch.float32) + a = torch.randn(1, total, HV, K, device=DEV, dtype=torch.float32) + b = torch.randn(1, total, HV, device=DEV, dtype=torch.float32) + A_log = torch.randn(HV, device=DEV, dtype=torch.float32) + dt_bias = torch.randn(HV, K, device=DEV, dtype=torch.float32) + + slots = torch.arange(1, bs + 1, device=DEV, dtype=torch.int32) + if pad_last and bs > 1: + slots[-1] = -1 + num_slots = bs + 1 + + h0_src = torch.randn(num_slots, HV, V, K, device=DEV, dtype=torch.float32) + max_len = int(lens.max()) + inter = torch.zeros(num_slots, max_len, HV, V, K, device=DEV, dtype=torch.float32) + + rawv = torch.zeros(num_slots, HV, L, V, device=DEV, dtype=torch.float32) + rawk = torch.zeros(num_slots, H, L, K, device=DEV, dtype=torch.float32) + gring = torch.zeros(num_slots, HV, L, K, device=DEV, dtype=torch.float32) + betar = torch.zeros(num_slots, HV, L, device=DEV, dtype=torch.float32) + + fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=h0_src, + initial_state_indices=slots, + cu_seqlens=cu, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + lower_bound=-5.0, + disable_state_update=True, + intermediate_states_buffer=inter, + intermediate_state_indices=slots, + cache_steps=max_len, + cache_ring=True, + replayssm_rawv=rawv, + replayssm_rawk=rawk, + replayssm_g=gring, + replayssm_beta=betar, + ) + + ckpt = h0_src.clone() + commit_kda_replayssm_spec( + ckpt, + rawv, + rawk, + gring, + betar, + slots, + acc, + max_cache_len=L, + num_k_heads=H, + use_qk_l2norm_in_kernel=True, + null_block_id=-1, + ) + + for j in range(bs): + if slots[j] < 0 or int(acc[j]) <= 0: + continue + base = inter[slots[j], int(acc[j]) - 1] + fold = ckpt[slots[j]] + rel = ((fold - base).abs().max() / base.abs().max().clamp_min(1e-6)).item() + assert ( + rel < 1e-3 + ), f"row={j} len={int(lens[j])} acc={int(acc[j])}: rel={rel:.3e}" + + +@pytest.mark.parametrize("bs,gamma,HV,H,K,V", SHAPES, ids=SHAPE_IDS) +def test_ragged_full_commit(bs, gamma, HV, H, K, V): + torch.manual_seed(0) + L = max(16, 2 * gamma) + lens = torch.randint(1, gamma + 1, (bs,), device=DEV, dtype=torch.int32) + _run_case(bs, gamma, HV, H, K, V, lens, acc=lens.clone(), L=L) + + +@pytest.mark.parametrize("bs,gamma,HV,H,K,V", SHAPES, ids=SHAPE_IDS) +def test_ragged_partial_commit(bs, gamma, HV, H, K, V): + torch.manual_seed(1) + L = max(16, 2 * gamma) + lens = torch.randint(1, gamma + 1, (bs,), device=DEV, dtype=torch.int32) + acc = (lens + 1) // 2 + _run_case(bs, gamma, HV, H, K, V, lens, acc=acc, L=L) + + +def test_ragged_pad_slot(): + torch.manual_seed(2) + bs, gamma, HV, H, K, V = 8, 8, 4, 4, 64, 64 + L = 16 + lens = torch.randint(1, gamma + 1, (bs,), device=DEV, dtype=torch.int32) + _run_case(bs, gamma, HV, H, K, V, lens, acc=lens.clone(), L=L, pad_last=True) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))