diff --git a/benchmark/kernels/bench_kda_verify_sweep.py b/benchmark/kernels/bench_kda_verify_sweep.py new file mode 100644 index 000000000..2990acd50 --- /dev/null +++ b/benchmark/kernels/bench_kda_verify_sweep.py @@ -0,0 +1,285 @@ +"""Sweep benchmark for the KDA chain-verify kernels (one layer, in-graph). + +Compares the four target-verify variants the KDA backend can dispatch: + + unfused causal_conv1d_update + recurrence, per-step ssm snapshots + unfused+ring causal_conv1d_update + recurrence, ReplaySSM CACHE_RING + fused fused_kda_conv_gating_verify, per-step ssm snapshots + fused+ring fused_kda_conv_gating_verify, ReplaySSM CACHE_RING + +Timing replays a CUDA graph capturing GRAPH_BATCH calls, matching how the +production verify runs (in-graph; bare launches would drown these ~10us +kernels in launch overhead). Imports only sglang.kernels.*, so it runs on +boxes where the sglang.srt/test import chain is broken. + + PYTHONPATH=python python3 benchmark/kernels/bench_kda_verify_sweep.py + ... --batch-sizes 1 4 16 64 --modes fused fused+ring + ... --sweep-bv # re-tune KDA_VERIFY_BLOCK_V per mode/batch + ... --hv-heads 16 # GQA shape (HV != H) +""" + +import argparse + +import torch + +import sglang.kernels.ops.attention.fla.fused_kda_conv_recurrent_verify as fused_mod +from sglang.kernels.ops.attention.fla.fused_kda_conv_recurrent_verify import ( + fused_kda_conv_gating_verify, +) +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update, +) +from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update + +_DEVICE = "cuda" +_DTYPE = torch.bfloat16 +_W = 4 +# Ring length: power of two >= 2 * draft tokens (memory_pool.py invariant). +_RING_LEN = 16 +_MODES = ("unfused", "unfused+ring", "fused", "fused+ring") +GRAPH_BATCH = 10 + + +def make_inputs(B, T, H, HV, K, V, seed=0): + torch.manual_seed(seed) + dim = 2 * H * K + HV * V + seq_len = B * T + lines = slots = B + 1 + rnd = lambda *s, dt=_DTYPE: torch.randn(*s, device=_DEVICE, dtype=dt) + return { + "mixed": rnd(seq_len, dim) * 0.5, + "w": rnd(dim, _W) * 0.3, + "bias": rnd(dim) * 0.1, + "a": rnd(seq_len, HV * K) * 0.5, + "b": rnd(seq_len, HV), + "A_log": rnd(HV, dt=torch.float32) * 0.5, + "dt_bias": rnd(HV * K, dt=torch.float32) * 0.5, + "conv_pool": rnd(lines, _W - 1, dim), + "ssm": rnd(slots, HV, V, K, dt=torch.float32) * 0.2, + "win_pool": torch.zeros(lines, T, _W - 1, dim, device=_DEVICE, dtype=_DTYPE), + "inter_ssm": torch.zeros( + lines, T, HV, V, K, device=_DEVICE, dtype=torch.float32 + ), + "rawv": rnd(slots, HV, _RING_LEN, V), + "rawk": rnd(slots, H, _RING_LEN, K), + "g": rnd(slots, HV, _RING_LEN, K, dt=torch.float32), + "beta": rnd(slots, HV, _RING_LEN, dt=torch.float32), + "cache_indices": torch.arange(B, device=_DEVICE, dtype=torch.int32), + "inter_indices": torch.arange(B, device=_DEVICE, dtype=torch.int32), + "cu": torch.arange(0, B + 1, device=_DEVICE, dtype=torch.int32) * T, + } + + +def _ring_kwargs(inp, on): + return dict( + cache_ring=on, + replayssm_rawv=inp["rawv"] if on else None, + replayssm_rawk=inp["rawk"] if on else None, + replayssm_g=inp["g"] if on else None, + replayssm_beta=inp["beta"] if on else None, + ) + + +def make_runner(mode, inp, B, T, H, HV, K, V, lower_bound=None): + dim = 2 * H * K + HV * V + seq_len = B * T + ring = mode.endswith("+ring") + scale = K**-0.5 + + if mode.startswith("fused"): + + def fn(): + fused_kda_conv_gating_verify( + mixed_qkv=inp["mixed"], + conv_weight=inp["w"], + conv_bias=inp["bias"], + conv_state=inp["conv_pool"].transpose(-1, -2), + conv_state_indices=inp["cache_indices"], + intermediate_conv_window=inp["win_pool"].transpose(-1, -2), + intermediate_state_indices=inp["inter_indices"], + a=inp["a"], + b=inp["b"], + A_log=inp["A_log"], + dt_bias=inp["dt_bias"], + ssm_states=inp["ssm"], + cache_indices=inp["cache_indices"], + intermediate_states_buffer=None if ring else inp["inter_ssm"], + scale=scale, + T=T, + num_q_heads=H, + num_v_heads=HV, + head_k_dim=K, + head_v_dim=V, + lower_bound=lower_bound, + **_ring_kwargs(inp, ring), + ) + + return fn + + def fn(): + x3 = inp["mixed"].reshape(B, T, dim).transpose(1, 2) + out3 = causal_conv1d_update( + x3, + inp["conv_pool"].transpose(-1, -2), + inp["w"], + inp["bias"], + activation="silu", + conv_state_indices=inp["cache_indices"], + intermediate_conv_window=inp["win_pool"].transpose(-1, -2), + intermediate_state_indices=inp["inter_indices"], + ) + mixed_out = out3.transpose(1, 2).reshape(seq_len, dim) + q, k, v = mixed_out.split([H * K, H * K, HV * V], dim=-1) + fused_sigmoid_gating_delta_rule_update( + A_log=inp["A_log"], + a=inp["a"], + dt_bias=inp["dt_bias"], + softplus_beta=1.0, + softplus_threshold=20.0, + q=q.unflatten(-1, (H, K)).unsqueeze(0), + k=k.unflatten(-1, (H, K)).unsqueeze(0), + v=v.unflatten(-1, (HV, V)).unsqueeze(0), + b=inp["b"], + initial_state_source=inp["ssm"], + initial_state_indices=inp["cache_indices"], + use_qk_l2norm_in_kernel=True, + cu_seqlens=inp["cu"], + is_kda=True, + disable_state_update=True, + intermediate_states_buffer=None if ring else inp["inter_ssm"], + intermediate_state_indices=None if ring else inp["inter_indices"], + cache_steps=T, + retrieve_parent_token=None, + lower_bound=lower_bound, + **_ring_kwargs(inp, ring), + ) + + return fn + + +def bench_graph(fn, iters=200): + """us per call, timed as CUDA-graph replays of GRAPH_BATCH captured calls.""" + for _ in range(3): # compile outside capture + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(GRAPH_BATCH): + fn() + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + graph.replay() + end.record() + torch.cuda.synchronize() + graph.reset() + return start.elapsed_time(end) * 1e3 / (iters * GRAPH_BATCH) + + +def check_ring_bitwise(B, T, H, HV, K, V, lower_bound=None): + """One-shot guard: fused+ring must fill the same ring bytes as unfused+ring.""" + ref, fus = (make_inputs(B, T, H, HV, K, V, seed=7) for _ in range(2)) + make_runner("unfused+ring", ref, B, T, H, HV, K, V, lower_bound)() + make_runner("fused+ring", fus, B, T, H, HV, K, V, lower_bound)() + torch.cuda.synchronize() + for name in ("rawv", "rawk", "g", "beta"): + assert torch.equal(ref[name], fus[name]), f"ring mismatch: {name}" + + +def run_modes(args, label_extra=""): + print( + f"H={args.heads} HV={args.hv_heads} K={args.head_k_dim} V={args.head_v_dim} " + f"T={args.draft_tokens} gate={'safe' if args.lower_bound is not None else 'std'} " + f"BV={fused_mod.KDA_VERIFY_BLOCK_V}{label_extra}" + ) + header = f"{'B':>4} " + "".join(f"{m:>14}" for m in args.modes) + print(header) + for B in args.batch_sizes: + times = [] + for mode in args.modes: + inp = make_inputs( + B, + args.draft_tokens, + args.heads, + args.hv_heads, + args.head_k_dim, + args.head_v_dim, + ) + fn = make_runner( + mode, + inp, + B, + args.draft_tokens, + args.heads, + args.hv_heads, + args.head_k_dim, + args.head_v_dim, + args.lower_bound, + ) + times.append(bench_graph(fn, iters=args.iters)) + row = f"{B:>4} " + "".join(f"{t:>11.2f} us" for t in times) + print(row) + print() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--heads", type=int, default=8) # ling-v3 TP4 KDA shape + parser.add_argument("--hv-heads", type=int, default=None) + parser.add_argument("--head-k-dim", type=int, default=128) + parser.add_argument("--head-v-dim", type=int, default=128) + parser.add_argument("--draft-tokens", type=int, default=4) + # ling-v3 runs the safe gate: --lower-bound -5.0 (kda_lower_bound). + parser.add_argument("--lower-bound", type=float, default=None) + parser.add_argument( + "--batch-sizes", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64] + ) + parser.add_argument("--modes", nargs="+", default=list(_MODES), choices=_MODES) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument( + "--sweep-bv", + action="store_true", + help="re-run the fused modes across KDA_VERIFY_BLOCK_V candidates; " + "BLOCK_V was tuned with snapshot writes on, so ring mode may move it", + ) + parser.add_argument("--skip-check", action="store_true") + args = parser.parse_args() + if args.hv_heads is None: + args.hv_heads = args.heads + if args.draft_tokens * 2 > _RING_LEN: + raise ValueError(f"--draft-tokens > {_RING_LEN // 2} exceeds the bench ring") + + if not args.skip_check: + check_ring_bitwise( + 4, + args.draft_tokens, + args.heads, + args.hv_heads, + args.head_k_dim, + args.head_v_dim, + args.lower_bound, + ) + print("ring bitwise check: OK\n") + + run_modes(args) + + if args.sweep_bv: + args.modes = [m for m in args.modes if m.startswith("fused")] or [ + "fused", + "fused+ring", + ] + default_bv = fused_mod.KDA_VERIFY_BLOCK_V + try: + for bv in (2, 4, 8, 16, 32): + fused_mod.KDA_VERIFY_BLOCK_V = bv + run_modes(args, label_extra=" (BV sweep)") + finally: + fused_mod.KDA_VERIFY_BLOCK_V = default_bv + + +if __name__ == "__main__": + main() diff --git a/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py b/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py index dfb5bc0ef..de00fa56b 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py +++ b/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py @@ -16,11 +16,17 @@ kernels. Requires ``T >= kernel_width - 1`` (the rolled conv state is then exactly the last ``kernel_width - 1`` input tokens, matching the reference kernel's store). -Numerics: deliberately bit-aligned with the unfused pair. The conv output is -rounded to the activation dtype (bf16) before entering the recurrence — -exactly what the unfused path does through its intermediate tensor — and all -expressions mirror the reference kernels line by line, with the same -num_warps so reduction order matches. +ReplaySSM (``cache_ring``): instead of per-step [HV, V, K] fp32 state +snapshots, stash each step's raw inputs (pre-l2norm k, pre-delta v, gate, +beta) into the per-slot rings the commit-time exact fold replays +(kda_replayssm_spec_decode.py) -- same CACHE_RING contract as the unfused +fused_sigmoid_gating_delta_rule_update, so the two paths fill identical rings. + +Numerics: aligned with the unfused pair. The conv output is rounded to the +activation dtype (bf16) before entering the recurrence — exactly what the +unfused path does through its intermediate tensor — and all expressions mirror +the reference kernels line by line. Reduction order still splits differently +where many V heads share one Q/K head, worth ~1 ulp on the output. """ from typing import Optional @@ -83,6 +89,18 @@ def fused_kda_conv_gating_verify_kernel( SAVE_INTERMEDIATE_WINDOW: tl.constexpr, CACHE_INTERMEDIATE_STATES: tl.constexpr, USE_GDC: tl.constexpr = False, + # ReplaySSM fused ring-write (spec verify): per-slot rings consumed by the + # commit-time exact fold (kda_replayssm_spec_decode.py). Off -> dead code. + replayssm_rawv=None, # [slots, HV, L, V] activation dtype + replayssm_rawk=None, # [slots, H, L, K] activation dtype + replayssm_g=None, # [slots, HV, L, K] fp32 + replayssm_beta=None, # [slots, HV, L] fp32 + stride_rawv_slot: tl.constexpr = 0, + stride_rawk_slot: tl.constexpr = 0, + stride_g_slot: tl.constexpr = 0, + stride_beta_slot: tl.constexpr = 0, + MAX_CACHE_LEN: tl.constexpr = 0, + CACHE_RING: tl.constexpr = False, ): # PDL: overlap prologue with the tail of the producer qkv-projection GEMM; # every global load (conv_state_indices, mixed_qkv, weights) happens after @@ -289,6 +307,53 @@ def fused_kda_conv_gating_verify_kernel( b_beta = 1.0 / (1.0 + tl.exp(-b_b)) + # ReplaySSM ring-write. Must sit here: b_k still pre-l2norm, b_v still + # pre-delta, b_g/b_beta formed -- so the commit fold's replay is + # bit-identical to the update below (mirrors the CACHE_RING block in + # fused_sigmoid_gating_recurrent.py). rawk dedups via is_qk_owner + # (per k-head); g/beta write once per v-head at i_v == 0. The + # t < MAX_CACHE_LEN guard drops absorb-overflow steps instead of + # smashing the next slot's ring. + if CACHE_RING: + if h0_idx >= 0 and t < MAX_CACHE_LEN: + ring_slot = h0_idx.to(tl.int64) + tl.store( + replayssm_rawv + + ring_slot * stride_rawv_slot + + i_hv * MAX_CACHE_LEN * V + + t * V + + o_v, + b_v.to(replayssm_rawv.dtype.element_ty), + mask=mask_v, + ) + if is_qk_owner: + tl.store( + replayssm_rawk + + ring_slot * stride_rawk_slot + + i_h * MAX_CACHE_LEN * K + + t * K + + o_k, + b_k.to(replayssm_rawk.dtype.element_ty), + mask=mask_k, + ) + if i_v == 0: + tl.store( + replayssm_g + + ring_slot * stride_g_slot + + i_hv * MAX_CACHE_LEN * K + + t * K + + o_k, + b_g, + mask=mask_k, + ) + tl.store( + replayssm_beta + + ring_slot * stride_beta_slot + + i_hv * MAX_CACHE_LEN + + t, + b_beta, + ) + if USE_QK_L2NORM_IN_KERNEL: b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6)) b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6)) @@ -358,15 +423,23 @@ def fused_kda_conv_gating_verify( softplus_beta: float = 1.0, softplus_threshold: float = 20.0, use_qk_l2norm_in_kernel: bool = True, - # num_warps=4 is ~1.3x faster than the unfused pair in-graph; the output, - # conv_state and conv-window caches stay bit-identical to the reference. - # Only the fp32 intermediate-ssm rollback cache differs: the tl.sum - # reduction-order delta (~1 ulp/step) compounds through the delta-rule - # recurrence — measured ~6e-8 at T=4 standard gate (the production MTP - # shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe gate. num_warps=1 - # reproduces the reference reduction order exactly (all buffers - # bit-identical) but is ~2.4x slower in-graph — numerics debugging only. + # num_warps=4 is ~1.3x faster than the unfused pair in-graph; conv_state + # and the conv-window cache stay bit-identical to the reference, the bf16 + # output within one ulp (the BV=4 tile reduces K in a different order). + # The fp32 intermediate-ssm rollback cache carries that ~1 ulp/step delta + # through the delta-rule recurrence — measured ~6e-8 at T=4 standard gate + # (the production MTP shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe + # gate. num_warps=1 is ~2.4x slower in-graph — numerics debugging only. + # The ReplaySSM ring values are bit-exact at any num_warps: they are + # elementwise (conv FMA chain, gate, sigmoid), upstream of every tl.sum. num_warps: int = 4, + # ReplaySSM fused ring-write; same parameter names as the unfused + # fused_sigmoid_gating_delta_rule_update so ring_kwargs pass through both. + cache_ring: bool = False, + replayssm_rawv: Optional[torch.Tensor] = None, + replayssm_rawk: Optional[torch.Tensor] = None, + replayssm_g: Optional[torch.Tensor] = None, + replayssm_beta: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Chain-verify fast path. Returns ``o`` of shape [1, seq_len, HV, V], matching the unfused ``target_verify`` output layout.""" @@ -420,6 +493,37 @@ def fused_kda_conv_gating_verify( if intermediate_states_buffer is not None: assert intermediate_states_buffer.is_contiguous() + if cache_ring: + # Per-layer ring views (memory_pool.py KDA spec rings). The kernel uses + # stride(0) as the slot pitch and packs within a slot from + # MAX_CACHE_LEN and the head/dim extents, so inner dims must be packed. + assert ( + replayssm_rawv is not None + and replayssm_rawk is not None + and replayssm_g is not None + and replayssm_beta is not None + ), "cache_ring requires all four replayssm_* rings" + max_cache_len = replayssm_rawv.shape[-2] + assert tuple(replayssm_rawv.shape[1:]) == (HV, max_cache_len, V) + assert tuple(replayssm_rawk.shape[1:]) == (H, max_cache_len, K) + assert tuple(replayssm_g.shape[1:]) == (HV, max_cache_len, K) + assert tuple(replayssm_beta.shape[1:]) == (HV, max_cache_len) + assert replayssm_rawv.stride()[1:] == (max_cache_len * V, V, 1) + assert replayssm_rawk.stride()[1:] == (max_cache_len * K, K, 1) + assert replayssm_g.stride()[1:] == (max_cache_len * K, K, 1) + assert replayssm_beta.stride()[1:] == (max_cache_len, 1) + assert replayssm_rawv.dtype == mixed_qkv.dtype + assert replayssm_rawk.dtype == mixed_qkv.dtype + assert replayssm_g.dtype == torch.float32 + assert replayssm_beta.dtype == torch.float32 + stride_rawv_slot = replayssm_rawv.stride(0) + stride_rawk_slot = replayssm_rawk.stride(0) + stride_g_slot = replayssm_g.stride(0) + stride_beta_slot = replayssm_beta.stride(0) + else: + max_cache_len = 0 + stride_rawv_slot = stride_rawk_slot = stride_g_slot = stride_beta_slot = 0 + grid = (NV, B * HV) # PDL (sm90+): chain behind the producer qkv-projection GEMM and signal the # downstream o_norm / o_proj. Scheduling only — bit-exactness unaffected. @@ -480,6 +584,16 @@ def fused_kda_conv_gating_verify( USE_LOWER_BOUND=lower_bound is not None, SAVE_INTERMEDIATE_WINDOW=intermediate_conv_window is not None, CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_g=replayssm_g, + replayssm_beta=replayssm_beta, + stride_rawv_slot=stride_rawv_slot, + stride_rawk_slot=stride_rawk_slot, + stride_g_slot=stride_g_slot, + stride_beta_slot=stride_beta_slot, + MAX_CACHE_LEN=max_cache_len, + CACHE_RING=cache_ring, # num_warps=1 matches the reference kernels' reduction order exactly; # higher values must be re-validated for bit-exactness before use. num_warps=num_warps, diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 40c728fab..b3b5cf8ab 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -39,6 +39,7 @@ from sglang.srt.runtime_context import ( get_disagg, get_exec, get_memory, + get_platform, get_spec, ) @@ -988,6 +989,27 @@ class KDAAttnBackend(MambaAttnBackendBase): "KDA target_verify requires a speculative mamba cache " "(MambaPool.SpeculativeState); none found." ) + # ReplaySSM: the ring-write is fused into the verify kernel + # (CACHE_RING) on both the fused chain-verify and the unfused triton + # paths; commit replays the ring instead of reading per-step + # snapshots. ring_kwargs stays empty for non-triton verify kernels, + # which never see replayssm. Ragged layouts work natively on the + # unfused path -- step_idx is the within-row step under varlen, so + # row i writes ring[slot][0..verify_lens[i]) and commit folds at most + # commit_lens of them (absorb overflow is bounded in-kernel). + replayssm_rawk = replayssm_g = replayssm_beta = None + ring_kwargs = {} + if replayssm_on: + replayssm_rawk = mamba_cache_params.replayssm_rawk + replayssm_g = mamba_cache_params.replayssm_g + replayssm_beta = mamba_cache_params.replayssm_beta + ring_kwargs = dict( + cache_ring=True, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_g=replayssm_g, + replayssm_beta=replayssm_beta, + ) intermediate_conv_window_cache = mamba_cache_params.intermediate_conv_window[0] intermediate_state_indices = self.verify_intermediate_state_indices @@ -1047,6 +1069,9 @@ class KDAAttnBackend(MambaAttnBackendBase): retrieve_next_sibling=retrieve_next_sibling, retrieve_parent_token=retrieve_parent_token, replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_g=replayssm_g, + replayssm_beta=replayssm_beta, ): return self._fused_chain_verify_fn( mixed_qkv=mixed_qkv, @@ -1077,6 +1102,7 @@ class KDAAttnBackend(MambaAttnBackendBase): head_k_dim=layer.head_k_dim, head_v_dim=layer.head_v_dim, lower_bound=layer.lower_bound, + **ring_kwargs, ) dense_token_indices = None mixed_qkv_dense = mixed_qkv.view(batch_size, draft_token_num, -1) @@ -1135,22 +1161,6 @@ class KDAAttnBackend(MambaAttnBackendBase): k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) - # ReplaySSM: the ring-write is fused into the triton verify kernel - # (CACHE_RING). Ragged layouts work natively -- step_idx is the - # within-row step under varlen, so row i writes - # ring[slot][0..verify_lens[i]) and commit folds at most commit_lens - # of them (absorb overflow is bounded in-kernel). ring_kwargs stays - # empty for non-triton verify kernels, which never see replayssm. - ring_kwargs = {} - if replayssm_rawv is not None: - ring_kwargs = dict( - cache_ring=True, - replayssm_rawv=replayssm_rawv, - replayssm_rawk=mamba_cache_params.replayssm_rawk, - replayssm_g=mamba_cache_params.replayssm_g, - replayssm_beta=mamba_cache_params.replayssm_beta, - ) - core_attn_out = self.kernel_dispatcher.target_verify( A_log=layer.A_log, dt_bias=layer.dt_bias, @@ -1210,10 +1220,13 @@ class KDAAttnBackend(MambaAttnBackendBase): retrieve_next_sibling: Optional[torch.Tensor], retrieve_parent_token: Optional[torch.Tensor], replayssm_rawv: Optional[torch.Tensor], + replayssm_rawk: Optional[torch.Tensor], + replayssm_g: Optional[torch.Tensor], + replayssm_beta: Optional[torch.Tensor], ) -> bool: if self._fused_chain_verify_fn is None or not mixed_qkv.is_cuda: return False - if replayssm_rawv is not None or any( + if any( value is not None for value in ( retrieve_next_token, @@ -1222,6 +1235,7 @@ class KDAAttnBackend(MambaAttnBackendBase): ) ): return False + replayssm_on = replayssm_rawv is not None if draft_token_num < 3 or mixed_qkv.shape[0] % draft_token_num != 0: return False if ( @@ -1240,6 +1254,15 @@ class KDAAttnBackend(MambaAttnBackendBase): seq_len, dim = mixed_qkv.shape batch_size = seq_len // draft_token_num + if replayssm_on and ( + batch_size != 1 or not (get_platform().is_sm90 or get_platform().is_sm100) + ): + # The runtime still uses BV=4, not the benchmark's best-BV sweep: + # fused+ring wins at B=1 but regresses from B=4 (B=2 at T=8) on + # both enabled architectures. Keep the ring path conservative until + # other batch/architecture combinations are measured. The snapshot + # path and the separate CuTe path are unchanged. + return False expected_dim = ( 2 * layer.num_q_heads * layer.head_k_dim + layer.num_v_heads * layer.head_v_dim @@ -1273,7 +1296,21 @@ class KDAAttnBackend(MambaAttnBackendBase): layer.A_log.dtype != torch.float32 or layer.dt_bias.dtype != torch.float32 or ssm_states.dtype != torch.float32 - or intermediate_state_cache is None + ): + return False + if replayssm_on: + if not self._replayssm_ring_ok( + layer=layer, + draft_token_num=draft_token_num, + mixed_qkv=mixed_qkv, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_g=replayssm_g, + replayssm_beta=replayssm_beta, + ): + return False + elif ( + intermediate_state_cache is None or intermediate_state_cache.dtype != torch.float32 ): return False @@ -1283,7 +1320,7 @@ class KDAAttnBackend(MambaAttnBackendBase): or b.stride(-1) != 1 or not conv_states.is_contiguous() or not ssm_states.is_contiguous() - or not intermediate_state_cache.is_contiguous() + or (not replayssm_on and not intermediate_state_cache.is_contiguous()) ): return False if ( @@ -1301,10 +1338,15 @@ class KDAAttnBackend(MambaAttnBackendBase): or ssm_states.ndim != 4 or tuple(ssm_states.shape[-3:]) != (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim) - or intermediate_state_cache.ndim != 5 - or intermediate_state_cache.shape[1] < draft_token_num - or tuple(intermediate_state_cache.shape[-3:]) - != (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim) + or ( + not replayssm_on + and ( + intermediate_state_cache.ndim != 5 + or intermediate_state_cache.shape[1] < draft_token_num + or tuple(intermediate_state_cache.shape[-3:]) + != (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim) + ) + ) ): return False if ( @@ -1324,15 +1366,76 @@ class KDAAttnBackend(MambaAttnBackendBase): b, conv_states, ssm_states, - intermediate_state_cache, intermediate_conv_window_cache, cache_indices, intermediate_state_indices, ) if layer.bias is not None: tensors += (layer.bias,) + # Ring devices are validated in _replayssm_ring_ok. + if not replayssm_on: + tensors += (intermediate_state_cache,) return all(tensor.device == mixed_qkv.device for tensor in tensors) + @staticmethod + def _replayssm_ring_ok( + *, + layer: RadixLinearAttention, + draft_token_num: int, + mixed_qkv: torch.Tensor, + replayssm_rawv: torch.Tensor, + replayssm_rawk: Optional[torch.Tensor], + replayssm_g: Optional[torch.Tensor], + replayssm_beta: Optional[torch.Tensor], + ) -> bool: + """Whether the per-layer ReplaySSM rings fit the fused ring-write. + + Layouts follow memory_pool.py's KDA spec rings: rawv [slots, HV, L, V] + and rawk [slots, H, L, K] in the activation dtype, g [slots, HV, L, K] + fp32 (per-K KDA gate), beta [slots, HV, L] fp32. The kernel uses + stride(0) as the slot pitch and assumes packed inner dims; anything + else falls back to the unfused path, which handles it. + """ + if replayssm_rawk is None or replayssm_g is None or replayssm_beta is None: + return False + if ( + replayssm_rawv.ndim != 4 + or replayssm_rawk.ndim != 4 + or replayssm_g.ndim != 4 + or replayssm_beta.ndim != 3 + ): + return False + H, HV = layer.num_q_heads, layer.num_v_heads + K, V = layer.head_k_dim, layer.head_v_dim + ring_len = replayssm_rawv.shape[-2] + if ring_len < draft_token_num: + return False + if ( + tuple(replayssm_rawv.shape[1:]) != (HV, ring_len, V) + or tuple(replayssm_rawk.shape[1:]) != (H, ring_len, K) + or tuple(replayssm_g.shape[1:]) != (HV, ring_len, K) + or tuple(replayssm_beta.shape[1:]) != (HV, ring_len) + ): + return False + if ( + replayssm_rawv.dtype != mixed_qkv.dtype + or replayssm_rawk.dtype != mixed_qkv.dtype + or replayssm_g.dtype != torch.float32 + or replayssm_beta.dtype != torch.float32 + ): + return False + if ( + replayssm_rawv.stride()[1:] != (ring_len * V, V, 1) + or replayssm_rawk.stride()[1:] != (ring_len * K, K, 1) + or replayssm_g.stride()[1:] != (ring_len * K, K, 1) + or replayssm_beta.stride()[1:] != (ring_len, 1) + ): + return False + return all( + ring.device == mixed_qkv.device + for ring in (replayssm_rawv, replayssm_rawk, replayssm_g, replayssm_beta) + ) + def _can_run_dspark_cutedsl_mtp( self, *, diff --git a/test/registered/kernel/attention/test_kda_fused_verify_backend.py b/test/registered/kernel/attention/test_kda_fused_verify_backend.py new file mode 100644 index 000000000..e042c12fb --- /dev/null +++ b/test/registered/kernel/attention/test_kda_fused_verify_backend.py @@ -0,0 +1,322 @@ +"""KDA backend dispatch and ReplaySSM verify -> commit -> verify parity.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from sglang.kernels.ops.attention.fla.fused_kda_conv_recurrent_verify import ( + fused_kda_conv_gating_verify, +) +from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( + HybridLinearAttnBackend, +) +from sglang.srt.layers.attention.linear.kda_backend import ( + KDAAttnBackend, + KDAKernelDispatcher, +) +from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend +from sglang.srt.mem_cache.memory_pool import MambaPool +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.runtime_context import override_platform +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +# One bf16 ulp: the fused and unfused kernels reduce K in different orders. +_OUTPUT_TOL = dict(rtol=2**-7, atol=1e-7) +# Fold and snapshot recurrence differ by ~1 fp32 ulp; a wrong-step commit +# moves elements by >= 1e-3, so 1e-5 still catches it. +_SNAPSHOT_ORACLE_TOL = dict(rtol=0, atol=1e-5) + + +class TestKDAFusedVerifyBackend(CustomTestCase): + def _make_case(self, batch_size=1, heads=2, v_heads=4, lower_bound=-5.0): + torch.manual_seed(36821) + steps, head_dim, num_layers = 4, 128, 2 + num_slots = batch_size + 3 + dim = (2 * heads + v_heads) * head_dim + + def randn(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, device="cuda", dtype=dtype) * 0.2 + + layers = [ + SimpleNamespace( + layer_id=i, + num_q_heads=heads, + num_k_heads=heads, + num_v_heads=v_heads, + head_q_dim=head_dim, + head_k_dim=head_dim, + head_v_dim=head_dim, + q_dim=heads * head_dim, + k_dim=heads * head_dim, + v_dim=v_heads * head_dim, + conv_weights=randn(dim, 4), + bias=randn(dim), + A_log=randn(v_heads, dtype=torch.float32), + dt_bias=randn(v_heads * head_dim, dtype=torch.float32), + lower_bound=lower_bound, + ) + for i in range(num_layers) + ] + state = MambaPool.SpeculativeState( + conv=[randn(num_layers, num_slots, 3, dim)], + temporal=randn( + num_layers, + num_slots, + v_heads, + head_dim, + head_dim, + dtype=torch.float32, + ), + intermediate_ssm=None, + intermediate_conv_window=[randn(num_layers, batch_size, steps, 3, dim)], + replayssm_rawv=randn(num_layers, num_slots, v_heads, 16, head_dim), + replayssm_rawk=randn(num_layers, num_slots, heads, 16, head_dim), + replayssm_g=randn( + num_layers, num_slots, v_heads, 16, head_dim, dtype=torch.float32 + ), + replayssm_beta=randn( + num_layers, num_slots, v_heads, 16, dtype=torch.float32 + ), + ) + # Physical slots differ from scratch rows; reverse them to catch callers + # accidentally committing by request index instead of mamba slot. + slots = torch.arange(batch_size + 1, 1, -1, device="cuda", dtype=torch.int32) + batch = SimpleNamespace( + forward_mode=ForwardMode.TARGET_VERIFY, + spec_info=SimpleNamespace(draft_token_num=steps, ragged_verify_layout=None), + ) + rounds = [ + [ + ( + randn(batch_size * steps, dim), + randn(1, batch_size * steps, v_heads * head_dim), + randn(1, batch_size * steps, v_heads), + ) + for _ in layers + ] + for _ in range(2) + ] + return layers, state, slots, batch, rounds + + def _make_backend(self, template, slots, steps, *, fused, ring=True): + state = MambaPool.SpeculativeState( + conv=[template.conv[0].clone()], + temporal=template.temporal.clone(), + intermediate_conv_window=[template.intermediate_conv_window[0].clone()], + intermediate_ssm=( + None + if ring + else template.temporal.new_zeros( + template.temporal.shape[0], + slots.numel(), + steps, + *template.temporal.shape[2:], + ) + ), + **{ + name: getattr(template, name).clone() if ring else None + for name in ( + "replayssm_rawv", + "replayssm_rawk", + "replayssm_g", + "replayssm_beta", + ) + }, + ) + # Only the model/pool setup is a fixture. Verify, dispatch, ring fold and + # conv rollback below all use the production backend and GPU kernels. + backend = KDAAttnBackend.__new__(KDAAttnBackend) + backend.req_to_token_pool = SimpleNamespace( + mamba2_layer_cache=state.at_layer_idx, + get_speculative_mamba2_params_all_layers=lambda: state, + mamba_pool=SimpleNamespace(replayssm_is_kda=ring), + ) + backend.forward_metadata = SimpleNamespace( + query_start_loc=torch.arange( + slots.numel() + 1, device="cuda", dtype=torch.int32 + ) + * steps, + mamba_cache_indices=slots, + retrieve_next_token=None, + retrieve_next_sibling=None, + retrieve_parent_token=None, + ) + backend.verify_intermediate_state_indices = torch.arange( + slots.numel(), device="cuda", dtype=torch.int32 + ) + backend.accept_lens_pool = None + backend.kernel_dispatcher = KDAKernelDispatcher( + LinearAttnKernelBackend.TRITON, + LinearAttnKernelBackend.TRITON, + LinearAttnKernelBackend.TRITON, + ) + backend._fused_chain_verify_fn = ( + Mock(wraps=fused_kda_conv_gating_verify) if fused else None + ) + hybrid = HybridLinearAttnBackend.__new__(HybridLinearAttnBackend) + hybrid.linear_attn_backend = backend + return backend, hybrid, state + + @staticmethod + def _verify(backend, layers, batch, inputs): + return [ + backend.forward_extend(layer, batch, mixed, a, b) + for layer, (mixed, a, b) in zip(layers, inputs) + ] + + def test_verify_commit_verify(self): + # B=1 exercises the enabled path. Platform override makes the dispatch + # testable on any CUDA CI runner; it does not replace a GPU kernel. + for platform, (heads, v_heads, lower_bound), num_accept_tokens in ( + ({"is_sm90": True}, (2, 2, None), 1), + ({"is_sm90": True}, (2, 2, None), 2), + ({"is_sm90": True}, (2, 2, None), 4), + ({"is_sm90": True}, (2, 4, -5.0), 1), + ({"is_sm90": True}, (2, 4, -5.0), 2), + ({"is_sm90": True}, (2, 4, -5.0), 4), + ({"is_sm90": False, "is_sm100": True}, (2, 4, -5.0), 2), + ): + with ( + self.subTest( + platform=platform, + heads=heads, + v_heads=v_heads, + lower_bound=lower_bound, + num_accept_tokens=num_accept_tokens, + ), + override_platform(**platform), + ): + layers, initial, slots, batch, rounds = self._make_case( + heads=heads, v_heads=v_heads, lower_bound=lower_bound + ) + fused, fused_hybrid, fused_state = self._make_backend( + initial, slots, 4, fused=True + ) + reference, ref_hybrid, ref_state = self._make_backend( + initial, slots, 4, fused=False + ) + snapshots, _, snapshot_state = self._make_backend( + initial, slots, 4, fused=False, ring=False + ) + out_fused = self._verify(fused, layers, batch, rounds[0]) + out_ref = self._verify(reference, layers, batch, rounds[0]) + self._verify(snapshots, layers, batch, rounds[0]) + for actual, expected in zip(out_fused, out_ref): + torch.testing.assert_close(actual, expected, **_OUTPUT_TOL) + for state in (fused_state, ref_state): + torch.testing.assert_close( + state.temporal, initial.temporal, rtol=0, atol=0 + ) + + last_steps = torch.full_like(slots, num_accept_tokens - 1) + for hybrid in (fused_hybrid, ref_hybrid): + hybrid.update_mamba_state_after_mtp_verify( + last_correct_step_indices=last_steps, + mamba_track_indices=None, + mamba_steps_to_track=None, + model=None, + ) + torch.testing.assert_close( + fused_state.temporal, ref_state.temporal, rtol=0, atol=0 + ) + torch.testing.assert_close( + fused_state.conv[0], ref_state.conv[0], rtol=0, atol=0 + ) + # Independent snapshot oracle: equality between two ring + # arms alone would miss a shared no-op / wrong-step commit. + expected_ssm = initial.temporal.clone() + expected_ssm[:, slots.long()] = snapshot_state.intermediate_ssm[ + :, :, num_accept_tokens - 1 + ] + torch.testing.assert_close( + fused_state.temporal, expected_ssm, **_SNAPSHOT_ORACLE_TOL + ) + expected_conv = initial.conv[0].clone() + for i, (mixed, _, _) in enumerate(rounds[0]): + history = torch.cat( + ( + initial.conv[0][i, slots.long()], + mixed.view(1, 4, -1)[:, :num_accept_tokens], + ), + dim=1, + ) + expected_conv[i, slots.long()] = history[:, -3:] + torch.testing.assert_close( + fused_state.conv[0], expected_conv, rtol=0, atol=0 + ) + + out_fused = self._verify(fused, layers, batch, rounds[1]) + out_ref = self._verify(reference, layers, batch, rounds[1]) + for actual, expected in zip(out_fused, out_ref): + torch.testing.assert_close(actual, expected, **_OUTPUT_TOL) + self.assertEqual(fused._fused_chain_verify_fn.call_count, 4) + + def test_ring_dispatch_falls_back(self): + # B=2 and the measured regression sizes on the enabled architectures, + # plus B=1 on an architecture without ring measurements. + sm90 = {"is_sm90": True, "is_sm100": False} + sm100 = {"is_sm90": False, "is_sm100": True} + other = {"is_sm90": False, "is_sm100": False} + for platform, batch_size in ( + (sm90, 2), + (sm90, 4), + (sm90, 16), + (sm90, 64), + (sm100, 2), + (sm100, 4), + (sm100, 16), + (other, 1), + (other, 4), + ): + with ( + self.subTest(platform=platform, batch_size=batch_size), + override_platform(**platform), + ): + layers, initial, slots, batch, rounds = self._make_case(batch_size) + backend, _, state = self._make_backend(initial, slots, 4, fused=True) + reference, _, ref_state = self._make_backend( + initial, slots, 4, fused=False + ) + out = self._verify(backend, layers, batch, rounds[0]) + ref = self._verify(reference, layers, batch, rounds[0]) + backend._fused_chain_verify_fn.assert_not_called() + for actual, expected in zip(out, ref): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for name in ( + "replayssm_rawv", + "replayssm_rawk", + "replayssm_g", + "replayssm_beta", + ): + torch.testing.assert_close( + getattr(state, name), getattr(ref_state, name), rtol=0, atol=0 + ) + + def test_snapshot_dispatch_is_unchanged(self): + for platform in ( + {"is_sm90": True, "is_sm100": False}, + {"is_sm90": False, "is_sm100": True}, + {"is_sm90": False, "is_sm100": False}, + ): + with self.subTest(platform=platform), override_platform(**platform): + layers, initial, slots, batch, rounds = self._make_case(batch_size=4) + backend, _, _ = self._make_backend( + initial, slots, 4, fused=True, ring=False + ) + reference, _, _ = self._make_backend( + initial, slots, 4, fused=False, ring=False + ) + out = self._verify(backend, layers, batch, rounds[0]) + ref = self._verify(reference, layers, batch, rounds[0]) + self.assertEqual(backend._fused_chain_verify_fn.call_count, 2) + for actual, expected in zip(out, ref): + torch.testing.assert_close(actual, expected, **_OUTPUT_TOL) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py b/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py index f67fe4634..0152d28b5 100644 --- a/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py +++ b/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py @@ -14,7 +14,7 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import ( ) from sglang.test.ci.ci_register import register_cuda_ci -register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=90, stage="base-b", runner_config="1-gpu-large") _DEVICE = "cuda" @@ -29,6 +29,35 @@ _CASES = [ (1, 4, 8, 8, 64, 64, 4, True, None, False, 8), ] +# ReplaySSM ring-write cases: _CASES plus HV != H shapes, which exercise the +# per-k-head rawk vs per-v-head g/beta writer split (the GQA hazard: a wrong +# head index scribbles another head's ring silently). +_RING_CASES = _CASES + [ + (2, 4, 2, 4, 128, 128, 4, True, None, False, 20), + (1, 5, 2, 8, 64, 64, 4, True, 1.5, False, 21), + (3, 4, 4, 8, 128, 128, 4, True, None, True, 22), +] + +# Power of two >= 2 * max draft T in _RING_CASES, matching the pool invariant +# (memory_pool.py: ring length must be a power of two >= 2 * draft tokens). +_RING_LEN = 16 + + +def _make_ring_buffers(H, HV, K, V): + # Garbage-filled so a full-tensor bitwise compare proves both that written + # positions match and that neither kernel scribbles outside them. + slots = 8 + return { + "rawv": torch.randn( + slots, HV, _RING_LEN, V, device=_DEVICE, dtype=torch.bfloat16 + ), + "rawk": torch.randn( + slots, H, _RING_LEN, K, device=_DEVICE, dtype=torch.bfloat16 + ), + "g": torch.randn(slots, HV, _RING_LEN, K, device=_DEVICE, dtype=torch.float32), + "beta": torch.randn(slots, HV, _RING_LEN, device=_DEVICE, dtype=torch.float32), + } + def _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed): torch.manual_seed(seed) @@ -71,13 +100,13 @@ def _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed): return inputs -def _run_reference(inp, B, T, H, HV, K, V, lower_bound): +def _run_reference(inp, B, T, H, HV, K, V, lower_bound, rings=None): dim = 2 * H * K + HV * V seq_len = B * T conv = inp["conv_pool"].clone() ssm = inp["ssm"].clone() win = inp["win_pool"].clone() - ic = inp["inter_ssm"].clone() + ic = inp["inter_ssm"].clone() if rings is None else None x3 = inp["mixed"].reshape(B, T, dim).transpose(1, 2) out3 = causal_conv1d_update( @@ -112,20 +141,27 @@ def _run_reference(inp, B, T, H, HV, K, V, lower_bound): cu_seqlens=cu, is_kda=True, disable_state_update=True, + # ReplaySSM mode (rings set) drops the per-step snapshots, exactly as + # the production GDN/KDA backends pass None + cache_ring. intermediate_states_buffer=ic, - intermediate_state_indices=inp["inter_indices"], + intermediate_state_indices=inp["inter_indices"] if rings is None else None, cache_steps=T, retrieve_parent_token=None, lower_bound=lower_bound, + cache_ring=rings is not None, + replayssm_rawv=rings["rawv"] if rings is not None else None, + replayssm_rawk=rings["rawk"] if rings is not None else None, + replayssm_g=rings["g"] if rings is not None else None, + replayssm_beta=rings["beta"] if rings is not None else None, ) return o, conv, win, ic -def _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps): +def _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps, rings=None): conv = inp["conv_pool"].clone() ssm = inp["ssm"].clone() win = inp["win_pool"].clone() - ic = inp["inter_ssm"].clone() + ic = inp["inter_ssm"].clone() if rings is None else None o = fused_kda_conv_gating_verify( mixed_qkv=inp["mixed"], @@ -150,18 +186,29 @@ def _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps): head_v_dim=V, lower_bound=lower_bound, num_warps=num_warps, + cache_ring=rings is not None, + replayssm_rawv=rings["rawv"] if rings is not None else None, + replayssm_rawk=rings["rawk"] if rings is not None else None, + replayssm_g=rings["g"] if rings is not None else None, + replayssm_beta=rings["beta"] if rings is not None else None, ) return o, conv, win, ic -def _compare_case(case, num_warps): +def _compare_case(case, num_warps, use_ring=False): B, T, H, HV, K, V, W, has_bias, lower_bound, neg_slot, seed = case inp = _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed) + if use_ring: + template = _make_ring_buffers(H, HV, K, V) + rings_ref = {name: buf.clone() for name, buf in template.items()} + rings_fus = {name: buf.clone() for name, buf in template.items()} + else: + rings_ref = rings_fus = None o_ref, conv_ref, win_ref, ic_ref = _run_reference( - inp, B, T, H, HV, K, V, lower_bound + inp, B, T, H, HV, K, V, lower_bound, rings=rings_ref ) o_fus, conv_fus, win_fus, ic_fus = _run_fused( - inp, B, T, H, HV, K, V, lower_bound, num_warps + inp, B, T, H, HV, K, V, lower_bound, num_warps, rings=rings_fus ) idx_vals = inp["idx_vals"] @@ -170,12 +217,21 @@ def _compare_case(case, num_warps): o_ref_v = o_ref.reshape(B, T, HV, V)[valid_rows] o_fus_v = o_fus.reshape(B, T, HV, V)[valid_rows] - assert torch.equal(o_ref_v, o_fus_v) + # One bf16 ulp: the fused and reference tiles reduce K in different orders. + torch.testing.assert_close(o_fus_v, o_ref_v, rtol=2**-7, atol=1e-7) assert torch.equal(conv_ref[touched_slots], conv_fus[touched_slots]) assert torch.equal(win_ref[valid_rows], win_fus[valid_rows]) - torch.testing.assert_close( - ic_ref[valid_rows], ic_fus[valid_rows], atol=4e-3, rtol=0 - ) + if use_ring: + # Full-tensor bitwise: ring values are elementwise (conv FMA chain, + # gate, sigmoid), upstream of every tl.sum, so they are exact at any + # num_warps; the shared garbage init makes any out-of-slot or + # negative-slot scribble a mismatch. + for name in ("rawv", "rawk", "g", "beta"): + assert torch.equal(rings_ref[name], rings_fus[name]), name + else: + torch.testing.assert_close( + ic_ref[valid_rows], ic_fus[valid_rows], atol=4e-3, rtol=0 + ) @pytest.mark.parametrize("case", _CASES) @@ -183,5 +239,10 @@ def test_matches_unfused_reference(case): _compare_case(case, num_warps=4) +@pytest.mark.parametrize("case", _RING_CASES) +def test_replayssm_ring_matches_unfused(case): + _compare_case(case, num_warps=4, use_ring=True) + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))