"""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()