219 lines
8.3 KiB
Python
219 lines
8.3 KiB
Python
"""Microbenchmark: buffered output-only GDN decode (ReplaySSM Part A) vs. the
|
|
existing packed GDN decode kernel.
|
|
|
|
Compares per-step decode latency of
|
|
``fused_recurrent_gated_delta_rule_packed_decode`` (writes the full recurrent
|
|
state S every step) against ``fused_recurrent_gdn_replayssm_decode`` at
|
|
L in {1, 8, 16} (writes the full state only every L steps) across batch sizes
|
|
{1, 16, 64, 256} for a realistic GDN config (HV=32, K=V=128).
|
|
|
|
The win is per-step HBM *state* traffic: the packed kernel reads + writes S
|
|
(~2 * num_slots * HV * V * K * 4 bytes / step for an fp32 state), while the
|
|
ReplaySSM kernel reads S every step but writes it only 1-in-L steps, plus a
|
|
small ring append (d:[HV,V], k:[H,K], g:[HV] per step). The amortized state
|
|
traffic ratio is reported per L.
|
|
|
|
Because the write happens only 1-in-L steps, the latency has to be amortized the
|
|
same way. The kernel branches on the caller's `write_pos` cursor and folds the
|
|
ring into the checkpoint only at `write_pos == L-1`, so both phases are timed and
|
|
combined: a cycle is L-1 non-flush steps and one flush step, and the flush step
|
|
costs roughly 2.3x a non-flush one. Timing a single pinned cursor position would
|
|
report whichever phase was pinned as if it were the per-step cost.
|
|
|
|
Run::
|
|
|
|
python benchmark/kernels/attention/bench_gdn_replayssm_decode.py
|
|
|
|
Requires a GPU (Triton).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
import triton
|
|
|
|
from sglang.kernels.ops.attention.fla.fused_recurrent import (
|
|
fused_recurrent_gated_delta_rule_packed_decode,
|
|
)
|
|
from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import (
|
|
fused_recurrent_gdn_replayssm_decode,
|
|
)
|
|
|
|
|
|
def _make_static(B, H, HV, K, V, dtype, device):
|
|
qk_dim = 2 * H * K
|
|
v_dim = HV * V
|
|
mixed_qkv = torch.randn(B, qk_dim + v_dim, device=device, dtype=dtype)
|
|
a = torch.randn(B, HV, device=device, dtype=dtype) * 0.5
|
|
b = torch.randn(B, HV, device=device, dtype=dtype)
|
|
A_log = (torch.randn(HV, device=device, dtype=torch.float32) * 0.3).contiguous()
|
|
dt_bias = (torch.randn(HV, device=device, dtype=torch.float32) * 0.1).contiguous()
|
|
return mixed_qkv, a, b, A_log, dt_bias
|
|
|
|
|
|
def _state_bytes_per_step(B, HV, K, V, L, dtype):
|
|
"""Amortized per-step HBM *state* traffic (bytes), state in fp32.
|
|
|
|
packed: read S + write S every step.
|
|
replay: read S every step; write S once per L steps; append ring records
|
|
(d:[HV,V] in `dtype`, k:[H,K] in `dtype` shared across HV//H, g:[HV]
|
|
fp32) every step. We report the dominant fp32-state terms; ring
|
|
appends are tiny by comparison and shown separately.
|
|
"""
|
|
fp32 = 4
|
|
state_elems = B * HV * V * K # one record per active request slot
|
|
packed = (state_elems * fp32) * 2 # read + write
|
|
replay = (state_elems * fp32) * (1 + 1.0 / L) # read every step + write 1/L
|
|
return packed, replay
|
|
|
|
|
|
def _amortize(t_nonflush, t_flush, L):
|
|
"""Mean per-step latency over one full L-phase cycle.
|
|
|
|
`write_pos` walks 0, 1, ... L-1 and wraps, so a steady-state cycle is L-1
|
|
non-flush steps and one flush step. The non-flush steps are not identical --
|
|
each replays one more ring entry than the last -- so the midpoint phase is
|
|
sampled as their representative rather than phase 0, which replays nothing
|
|
and is the cheapest step in the cycle. Against a full per-phase sweep this
|
|
is accurate to <1%, where using phase 0 understates the mean by ~3%.
|
|
"""
|
|
if L == 1:
|
|
return t_flush
|
|
return ((L - 1) * t_nonflush + t_flush) / L
|
|
|
|
|
|
def _bench_cfg(
|
|
B, H, HV, K, V, Ls, dtype, device, nk=2, num_slots=None, warmup=25, rep=100
|
|
):
|
|
num_slots = num_slots or B
|
|
mixed_qkv, a, b, A_log, dt_bias = _make_static(B, H, HV, K, V, dtype, device)
|
|
scale = K**-0.5
|
|
cache_indices = torch.arange(B, device=device, dtype=torch.int32)
|
|
|
|
# packed decode
|
|
state = torch.randn(num_slots, HV, V, K, device=device, dtype=torch.float32)
|
|
out = mixed_qkv.new_empty(B, 1, HV, V)
|
|
|
|
def run_packed():
|
|
fused_recurrent_gated_delta_rule_packed_decode(
|
|
mixed_qkv=mixed_qkv,
|
|
a=a,
|
|
b=b,
|
|
A_log=A_log,
|
|
dt_bias=dt_bias,
|
|
scale=scale,
|
|
initial_state=state,
|
|
out=out,
|
|
ssm_state_indices=cache_indices,
|
|
use_qk_l2norm_in_kernel=True,
|
|
)
|
|
|
|
t_packed = triton.testing.do_bench(run_packed, warmup=warmup, rep=rep)
|
|
|
|
rows = []
|
|
for L in Ls:
|
|
rstate = torch.randn(num_slots, HV, V, K, device=device, dtype=torch.float32)
|
|
d_cache = torch.zeros(num_slots, HV, L, V, device=device, dtype=dtype)
|
|
k_cache = torch.zeros(num_slots, H, L, K, device=device, dtype=dtype)
|
|
g_cache = torch.zeros(num_slots, HV, L, device=device, dtype=torch.float32)
|
|
write_pos = torch.empty(B, device=device, dtype=torch.int32)
|
|
rout = mixed_qkv.new_empty(B, 1, HV, V)
|
|
|
|
def run_replay():
|
|
fused_recurrent_gdn_replayssm_decode(
|
|
mixed_qkv=mixed_qkv,
|
|
a=a,
|
|
b=b,
|
|
A_log=A_log,
|
|
dt_bias=dt_bias,
|
|
scale=scale,
|
|
initial_state=rstate,
|
|
d_cache=d_cache,
|
|
k_cache=k_cache,
|
|
g_cache=g_cache,
|
|
out=rout,
|
|
ssm_state_indices=cache_indices,
|
|
write_pos=write_pos,
|
|
use_qk_l2norm_in_kernel=True,
|
|
nk=nk,
|
|
)
|
|
|
|
def time_at(pos):
|
|
write_pos.fill_(pos)
|
|
return triton.testing.do_bench(run_replay, warmup=warmup, rep=rep)
|
|
|
|
# The kernel branches on write_pos: it appends to the ring on every step
|
|
# but folds the ring into the checkpoint only when write_pos == L-1. Both
|
|
# phases have to be timed, or the reported cost is the cheap one repeated
|
|
# (see `_amortize`). L == 1 is all flush, so it has no non-flush phase.
|
|
t_flush = time_at(L - 1)
|
|
t_nonflush = time_at((L - 1) // 2) if L > 1 else None
|
|
t_replay = _amortize(t_nonflush, t_flush, L)
|
|
|
|
packed_bytes, replay_bytes = _state_bytes_per_step(B, HV, K, V, L, dtype)
|
|
rows.append(
|
|
(
|
|
L,
|
|
t_replay,
|
|
t_nonflush,
|
|
t_flush,
|
|
t_packed / t_replay,
|
|
replay_bytes / packed_bytes,
|
|
)
|
|
)
|
|
return t_packed, rows
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--hv", type=int, default=32, help="num value heads")
|
|
parser.add_argument("--h", type=int, default=16, help="num key/query heads")
|
|
parser.add_argument("--k", type=int, default=128)
|
|
parser.add_argument("--v", type=int, default=128)
|
|
parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 16, 64, 256])
|
|
parser.add_argument("--ls", type=int, nargs="+", default=[1, 8, 16])
|
|
parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16")
|
|
parser.add_argument(
|
|
"--nk", type=int, default=2, help="K-chunks the reconstruction walks"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if not torch.cuda.is_available():
|
|
raise SystemExit("CUDA / Triton required for this microbenchmark.")
|
|
device = "cuda"
|
|
dtype = {
|
|
"bf16": torch.bfloat16,
|
|
"fp16": torch.float16,
|
|
"fp32": torch.float32,
|
|
}[args.dtype]
|
|
|
|
print(
|
|
f"GDN ReplaySSM decode microbench HV={args.hv} H={args.h} "
|
|
f"K={args.k} V={args.v} dtype={args.dtype} nk={args.nk}\n"
|
|
"replay = mean per-step latency over one L-phase cycle (ms), split into "
|
|
"its\nnon-flush and flush steps; speedup = packed/replay; "
|
|
"state-traffic = replay/packed\n(lower is better)"
|
|
)
|
|
for B in args.batch_sizes:
|
|
t_packed, rows = _bench_cfg(
|
|
B, args.h, args.hv, args.k, args.v, args.ls, dtype, device, nk=args.nk
|
|
)
|
|
print(f"\nB={B:<4d} packed={t_packed:.4f} ms")
|
|
for L, t_replay, t_nonflush, t_flush, speedup, traffic_ratio in rows:
|
|
split = (
|
|
f"(flush only)"
|
|
if t_nonflush is None
|
|
else f"(non-flush {t_nonflush:.4f}, flush {t_flush:.4f})"
|
|
)
|
|
print(
|
|
f" L={L:<3d} replay={t_replay:.4f} ms "
|
|
f"speedup={speedup:5.2f}x "
|
|
f"state-traffic={traffic_ratio:5.2f}x {split}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|