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

Add kda replayssm tests (#33630)

Browse Source
This commit is contained in:
Ke Bao
2026-08-10 09:11:20 +08:00
committed by GitHub
parent 25b7015064
commit 68b961e9fb
5 changed files with 776 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/test_kda_mtp_cutedsl_replayssm_ring.py
+194
View File
@@ -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"]))
test/registered/kernels/test_kda_replayssm_fold.py
+147
View File
@@ -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()
test/registered/kernels/test_kda_replayssm_fold_batched.py
+140
View File
@@ -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"]))
test/registered/kernels/test_kda_replayssm_ring_fused.py
+145
View File
@@ -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<gamma, single head, and padding (-1) slots. Both gate branches
(safe gate / softplus). A wrong index shows up as a shape-specific mismatch.
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 fused 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, T, HV, H, K, V) -- cover GQA (H<HV), non-pow2 K/V, small T, single head.
SHAPES = [
(16, 7, 4, 4, 64, 64), # baseline square heads
(64, 7, 32, 32, 128, 128), # K3-like TP8 shape
(8, 4, 8, 2, 128, 128), # GQA: 4 v-heads per k-head
(4, 3, 6, 3, 96, 80), # non-pow2 K/V, GQA 2:1
(1, 1, 4, 4, 64, 64), # single req, single step
]
SHAPE_IDS = ["square", "k3-tp8", "gqa4", "nonpow2", "single"]
GATES = [(-5.0, "safe"), (None, "softplus")]
@pytest.mark.parametrize("bs,T,HV,H,K,V", SHAPES, ids=SHAPE_IDS)
@pytest.mark.parametrize(
"lower_bound", [g[0] for g in GATES], ids=[g[1] for g in GATES]
)
@pytest.mark.parametrize("pad", [False, True], ids=["nopad", "pad"])
def test_ring_fold_parity(bs, T, HV, H, K, V, lower_bound, pad):
L = max(16, 2 * T) # ring length; power-of-two backstop satisfied
scale = K**-0.5
torch.manual_seed(0)
q = torch.randn(bs, T, H, K, device=DEV, dtype=torch.float32)
k = torch.randn(bs, T, H, K, device=DEV, dtype=torch.float32)
v = torch.randn(bs, T, HV, V, device=DEV, dtype=torch.float32)
a = torch.randn(bs, T, HV, K, device=DEV, dtype=torch.float32)
b = torch.randn(bs, 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(bs, HV, V, K, device=DEV, dtype=torch.float32)
# slots 1..bs; optionally set one row to a -1 padding slot (must be skipped).
slots = torch.arange(1, bs + 1, device=DEV, dtype=torch.int32)
if pad and bs > 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"]))
test/registered/kernels/test_kda_replayssm_ring_ragged.py
+150
View File
@@ -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"]))