minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[KDA] Support ReplaySSM ring-write in the fused chain-verify kernel (#36821)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-09-15 10:22:41 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent a23fd557ed
commit 99060191e7
5 changed files with 935 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/bench_kda_verify_sweep.py
+285
View File
@@ -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()
python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py
+127 -13
View File
@@ -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,
python/sglang/srt/layers/attention/linear/kda_backend.py
+127 -24
View File
@@ -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,
*,
test/registered/kernel/attention/test_kda_fused_verify_backend.py
+322
View File
@@ -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()
test/registered/kernels/test_fused_kda_conv_recurrent_verify.py
+74 -13
View File
@@ -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__]))