chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
481
third_party/sglang/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py
vendored
Normal file
481
third_party/sglang/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py
vendored
Normal file
@@ -0,0 +1,481 @@
|
||||
"""Benchmark & Correctness: CuTe DSL KDA Decode vs Triton KDA Decode.
|
||||
|
||||
This benchmark assumes the production / Triton canonical state layout:
|
||||
ssm_states.shape == (pool_size, HV, V, K)
|
||||
|
||||
Both the Triton baseline and the CuTe DSL candidate operate directly on that VK
|
||||
layout. No transpose is performed anywhere in the benchmark.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python"))
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.cutedsl_kda import cutedsl_fused_sigmoid_gating_kda_update
|
||||
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
|
||||
fused_sigmoid_gating_delta_rule_update,
|
||||
)
|
||||
from sglang.srt.layers.attention.fla.kda import chunk_kda
|
||||
|
||||
|
||||
def make_inputs(
|
||||
B: int,
|
||||
H: int,
|
||||
HV: int,
|
||||
K: int,
|
||||
V: int,
|
||||
pool_size: int,
|
||||
device: str,
|
||||
dtype: torch.dtype,
|
||||
layout: str,
|
||||
seed: int = 42,
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
|
||||
assert K == 128
|
||||
assert V % 16 == 0 and V % 32 == 0
|
||||
|
||||
if layout == "varlen":
|
||||
q = torch.randn(1, B, H, K, device=device, dtype=dtype)
|
||||
k = torch.randn(1, B, H, K, device=device, dtype=dtype)
|
||||
v = torch.randn(1, B, HV, V, device=device, dtype=dtype)
|
||||
|
||||
# decode params
|
||||
a = torch.randn(B, HV, K, device=device, dtype=dtype)
|
||||
b = torch.randn(B, HV, device=device, dtype=dtype)
|
||||
|
||||
# prefill params for chunk_kda must keep batch dim = 1
|
||||
# chunk_kda requires g, beta, v to have the same head count as k (H),
|
||||
# matching the real KimiLinear model where num_heads == num_kv_heads.
|
||||
prefill_v = torch.randn(1, B, H, V, device=device, dtype=dtype)
|
||||
prefill_g = torch.randn(1, B, H, K, device=device, dtype=dtype)
|
||||
prefill_beta = torch.sigmoid(torch.randn(1, B, H, device=device, dtype=dtype))
|
||||
|
||||
cu_seqlens = torch.arange(B + 1, device=device, dtype=torch.int32)
|
||||
|
||||
elif layout == "dense":
|
||||
q = torch.randn(B, 1, H, K, device=device, dtype=dtype)
|
||||
k = torch.randn(B, 1, H, K, device=device, dtype=dtype)
|
||||
v = torch.randn(B, 1, HV, V, device=device, dtype=dtype)
|
||||
|
||||
# decode params
|
||||
a = torch.randn(B, 1, HV, K, device=device, dtype=dtype)
|
||||
b = torch.randn(B, 1, HV, device=device, dtype=dtype)
|
||||
|
||||
# prefill params for chunk_kda dense path
|
||||
# chunk_kda requires g, beta, v to have the same head count as k (H),
|
||||
# matching the real KimiLinear model where num_heads == num_kv_heads.
|
||||
prefill_v = torch.randn(B, 1, H, V, device=device, dtype=dtype)
|
||||
prefill_g = torch.randn(B, 1, H, K, device=device, dtype=dtype)
|
||||
prefill_beta = torch.sigmoid(torch.randn(B, 1, H, device=device, dtype=dtype))
|
||||
|
||||
cu_seqlens = torch.arange(B + 1, device=device, dtype=torch.int32)
|
||||
else:
|
||||
raise ValueError(f"Unknown layout: {layout}")
|
||||
|
||||
A_log = torch.randn(HV, device=device, dtype=torch.float32)
|
||||
dt_bias = torch.randn(HV, K, device=device, dtype=dtype)
|
||||
|
||||
ssm_states = (
|
||||
torch.randn(pool_size, HV, V, K, device=device, dtype=torch.float32) * 0.1
|
||||
)
|
||||
cache_indices = torch.arange(B, device=device, dtype=torch.int32)
|
||||
|
||||
return dict(
|
||||
B=B,
|
||||
H=H,
|
||||
HV=HV,
|
||||
K=K,
|
||||
V=V,
|
||||
pool_size=pool_size,
|
||||
layout=layout,
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
a=a,
|
||||
b=b,
|
||||
prefill_v=prefill_v,
|
||||
prefill_g=prefill_g,
|
||||
prefill_beta=prefill_beta,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
ssm_states=ssm_states,
|
||||
cache_indices=cache_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
)
|
||||
|
||||
|
||||
def run_baseline(inp):
|
||||
state = inp["ssm_states"].clone()
|
||||
o = fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=inp["v"],
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=state,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
is_kda=True,
|
||||
)
|
||||
return o, state
|
||||
|
||||
|
||||
def run_cutedsl(inp):
|
||||
state = inp["ssm_states"].clone()
|
||||
o = cutedsl_fused_sigmoid_gating_kda_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=inp["v"],
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=state,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
)
|
||||
return o, state
|
||||
|
||||
|
||||
def run_prefill_then_decode_baseline(inp):
|
||||
ssm_states = inp["ssm_states"].clone()
|
||||
prefill_v_clone = inp["prefill_v"].clone()
|
||||
v_clone = inp["v"].clone()
|
||||
|
||||
_ = chunk_kda(
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=prefill_v_clone,
|
||||
g=inp["prefill_g"],
|
||||
beta=inp["prefill_beta"],
|
||||
initial_state=ssm_states,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
)
|
||||
|
||||
o = fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=v_clone,
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=ssm_states,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
is_kda=True,
|
||||
)
|
||||
return o, ssm_states
|
||||
|
||||
|
||||
def run_prefill_then_decode_cutedsl(inp):
|
||||
ssm_states = inp["ssm_states"].clone()
|
||||
prefill_v_clone = inp["prefill_v"].clone()
|
||||
v_clone = inp["v"].clone()
|
||||
|
||||
_ = chunk_kda(
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=prefill_v_clone,
|
||||
g=inp["prefill_g"],
|
||||
beta=inp["prefill_beta"],
|
||||
initial_state=ssm_states,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
)
|
||||
|
||||
o = cutedsl_fused_sigmoid_gating_kda_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=v_clone,
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=ssm_states,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
)
|
||||
return o, ssm_states
|
||||
|
||||
|
||||
def _assert_close(name, x, y, atol=3e-2, rtol=2e-2):
|
||||
try:
|
||||
torch.testing.assert_close(x.float(), y.float(), atol=atol, rtol=rtol)
|
||||
return True, 0.0
|
||||
except AssertionError:
|
||||
max_diff = (x - y).abs().max().item()
|
||||
return False, max_diff
|
||||
|
||||
|
||||
def check_correctness(B, H, HV, K, V, pool_size, device, dtype, layout):
|
||||
tag = (
|
||||
f"layout={layout:<6} B={B:>4} H={H:>2} HV={HV:>2} "
|
||||
f"K={K:>3} V={V:>3} pool={pool_size:>4}"
|
||||
)
|
||||
inp = make_inputs(B, H, HV, K, V, pool_size, device, dtype, layout)
|
||||
|
||||
o_ref, st_ref = run_baseline(inp)
|
||||
o_cute, st_cute = run_cutedsl(inp)
|
||||
|
||||
ok_o, diff_o = _assert_close("output", o_cute, o_ref)
|
||||
valid_mask = inp["cache_indices"] >= 0
|
||||
valid_idx = inp["cache_indices"][valid_mask]
|
||||
ok_s, diff_s = _assert_close("state", st_cute[valid_idx], st_ref[valid_idx])
|
||||
|
||||
if ok_o and ok_s:
|
||||
print(f" [PASS] {tag}")
|
||||
return True
|
||||
|
||||
details = []
|
||||
if not ok_o:
|
||||
details.append(f"output max_diff={diff_o:.6f}")
|
||||
if not ok_s:
|
||||
details.append(f"state max_diff={diff_s:.6f}")
|
||||
print(f" [FAIL] {tag} ({', '.join(details)})")
|
||||
return False
|
||||
|
||||
|
||||
def check_prefill_chain(B, H, HV, K, V, pool_size, device, dtype, layout):
|
||||
tag = (
|
||||
f"[prefill->decode] layout={layout:<6} B={B:>4} H={H:>2} HV={HV:>2} "
|
||||
f"K={K:>3} V={V:>3} pool={pool_size:>4}"
|
||||
)
|
||||
inp = make_inputs(B, H, HV, K, V, pool_size, device, dtype, layout)
|
||||
|
||||
o_ref, st_ref = run_prefill_then_decode_baseline(inp)
|
||||
o_cute, st_cute = run_prefill_then_decode_cutedsl(inp)
|
||||
|
||||
ok_o, diff_o = _assert_close("output", o_cute, o_ref)
|
||||
valid_mask = inp["cache_indices"] >= 0
|
||||
valid_idx = inp["cache_indices"][valid_mask]
|
||||
ok_s, diff_s = _assert_close("state", st_cute[valid_idx], st_ref[valid_idx])
|
||||
|
||||
if ok_o and ok_s:
|
||||
print(f" [PASS] {tag}")
|
||||
return True
|
||||
|
||||
details = []
|
||||
if not ok_o:
|
||||
details.append(f"output max_diff={diff_o:.6f}")
|
||||
if not ok_s:
|
||||
details.append(f"state max_diff={diff_s:.6f}")
|
||||
print(f" [FAIL] {tag} ({', '.join(details)})")
|
||||
return False
|
||||
|
||||
|
||||
def bench_shape(B, H, HV, K, V, pool_size, device, dtype, layout):
|
||||
inp = make_inputs(B, H, HV, K, V, pool_size, device, dtype, layout)
|
||||
|
||||
def fn_triton():
|
||||
fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=inp["v"],
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=inp["ssm_states"],
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
is_kda=True,
|
||||
)
|
||||
|
||||
def fn_cute():
|
||||
cutedsl_fused_sigmoid_gating_kda_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=inp["v"],
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=inp["ssm_states"],
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"] if inp["layout"] == "varlen" else None,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
)
|
||||
|
||||
for _ in range(10):
|
||||
fn_triton()
|
||||
fn_cute()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
try:
|
||||
ms_triton, _, _ = triton.testing.do_bench(
|
||||
fn_triton, quantiles=[0.5, 0.2, 0.8], warmup=50, rep=200
|
||||
)
|
||||
ms_cute, _, _ = triton.testing.do_bench(
|
||||
fn_cute, quantiles=[0.5, 0.2, 0.8], warmup=50, rep=200
|
||||
)
|
||||
except Exception:
|
||||
rep = 100
|
||||
st = time.perf_counter()
|
||||
for _ in range(rep):
|
||||
fn_triton()
|
||||
torch.cuda.synchronize()
|
||||
ms_triton = (time.perf_counter() - st) / rep * 1000
|
||||
|
||||
st = time.perf_counter()
|
||||
for _ in range(rep):
|
||||
fn_cute()
|
||||
torch.cuda.synchronize()
|
||||
ms_cute = (time.perf_counter() - st) / rep * 1000
|
||||
|
||||
speedup = ms_triton / ms_cute if ms_cute > 0 else float("inf")
|
||||
delta = (ms_cute - ms_triton) * 1000
|
||||
print(
|
||||
f" {layout:>6} {B:>5} {H:>3} {HV:>3} {K:>3} {V:>3} | "
|
||||
f"{ms_triton * 1000:>12.1f} | "
|
||||
f"{ms_cute * 1000:>13.1f} | "
|
||||
f"{speedup:>8.2f} | "
|
||||
f"{delta:>11.1f}"
|
||||
)
|
||||
|
||||
|
||||
def run_correctness(device, dtype):
|
||||
print("=" * 78)
|
||||
print("Correctness: Triton KDA Decode vs CuTe DSL KDA Decode")
|
||||
print("=" * 78)
|
||||
|
||||
shapes = [
|
||||
("dense", 1, 8, 16, 128, 128, 32),
|
||||
("dense", 4, 8, 16, 128, 128, 32),
|
||||
("dense", 32, 8, 16, 128, 128, 128),
|
||||
("dense", 64, 8, 16, 128, 128, 128),
|
||||
("varlen", 4, 8, 16, 128, 128, 32),
|
||||
("varlen", 16, 8, 16, 128, 128, 64),
|
||||
("varlen", 32, 8, 16, 128, 128, 128),
|
||||
("varlen", 64, 8, 16, 128, 128, 128),
|
||||
("varlen", 1, 16, 32, 128, 128, 32),
|
||||
("varlen", 32, 16, 32, 128, 128, 128),
|
||||
("varlen", 64, 16, 16, 128, 128, 128),
|
||||
]
|
||||
|
||||
all_pass = True
|
||||
for layout, B, H, HV, K, V, pool_size in shapes:
|
||||
if not check_correctness(B, H, HV, K, V, pool_size, device, dtype, layout):
|
||||
all_pass = False
|
||||
|
||||
print()
|
||||
print("=" * 78)
|
||||
print("Correctness: Triton prefill/extend -> CuTe decode chain")
|
||||
print("=" * 78)
|
||||
for layout, B, H, HV, K, V, pool_size in shapes[:8]:
|
||||
if not check_prefill_chain(B, H, HV, K, V, pool_size, device, dtype, layout):
|
||||
all_pass = False
|
||||
|
||||
print()
|
||||
print("ALL PASSED." if all_pass else "SOME FAILED.")
|
||||
return all_pass
|
||||
|
||||
|
||||
def run_benchmark(device, dtype):
|
||||
print()
|
||||
print("=" * 92)
|
||||
print("Benchmark: Triton KDA Decode vs CuTe DSL KDA Decode")
|
||||
print("=" * 92)
|
||||
|
||||
bench_configs = [
|
||||
("dense", 1, 8, 16),
|
||||
("dense", 4, 8, 16),
|
||||
("dense", 32, 8, 16),
|
||||
("dense", 64, 8, 16),
|
||||
("varlen", 1, 8, 16),
|
||||
("varlen", 4, 8, 16),
|
||||
("varlen", 8, 8, 16),
|
||||
("varlen", 16, 8, 16),
|
||||
("varlen", 32, 8, 16),
|
||||
("varlen", 64, 8, 16),
|
||||
("varlen", 128, 8, 16),
|
||||
("varlen", 32, 16, 32),
|
||||
("varlen", 64, 16, 16),
|
||||
]
|
||||
|
||||
K = 128
|
||||
V = 128
|
||||
pool_size = 512
|
||||
|
||||
print(f" Config: K={K}, V={V}, pool_size={pool_size}, dtype={dtype}")
|
||||
print(
|
||||
f" {'layout':>6} {'B':>5} {'H':>3} {'HV':>3} {'K':>3} {'V':>3} | "
|
||||
f"{'triton (μs)':>12} | "
|
||||
f"{'cutedsl (μs)':>13} | "
|
||||
f"{'speedup':>8} | "
|
||||
f"{'delta (μs)':>11}"
|
||||
)
|
||||
print(" " + "-" * 82)
|
||||
|
||||
for layout, B, H, HV in bench_configs:
|
||||
actual_pool = max(pool_size, B + 16)
|
||||
bench_shape(B, H, HV, K, V, actual_pool, device, dtype, layout)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark & Correctness: Triton KDA Decode vs CuTe DSL KDA Decode"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=["all", "correctness", "bench"],
|
||||
default="all",
|
||||
help="Run mode (default: all)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
choices=["float16", "bfloat16", "float32"],
|
||||
default="bfloat16",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
device = "cuda"
|
||||
dtype = getattr(torch, args.dtype)
|
||||
|
||||
cap = torch.cuda.get_device_capability()
|
||||
dev_name = torch.cuda.get_device_name()
|
||||
print(f"Device: {dev_name} (SM {cap[0]}{cap[1]})")
|
||||
|
||||
if args.mode in ("all", "correctness"):
|
||||
all_pass = run_correctness(device, dtype)
|
||||
if not all_pass and args.mode == "all":
|
||||
print("\nSkipping benchmark due to correctness failures.")
|
||||
return 1
|
||||
|
||||
if args.mode in ("all", "bench"):
|
||||
run_benchmark(device, dtype)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
488
third_party/sglang/benchmark/bench_linear_attention/bench_gdn_decode.py
vendored
Normal file
488
third_party/sglang/benchmark/bench_linear_attention/bench_gdn_decode.py
vendored
Normal file
@@ -0,0 +1,488 @@
|
||||
"""
|
||||
Benchmark & Correctness: GDN Packed Decode vs Baseline Decode.
|
||||
|
||||
Compares:
|
||||
- Baseline: split(mixed_qkv) → view → fused_sigmoid_gating_delta_rule_update
|
||||
- Packed: fused_recurrent_gated_delta_rule_packed_decode (single kernel)
|
||||
|
||||
The packed path eliminates:
|
||||
- torch.split() + .view() tensor materialization
|
||||
- Separate gating kernel launches
|
||||
- Intermediate tensor allocations
|
||||
|
||||
Reports correctness (output & state matching) and performance (ms, speedup).
|
||||
|
||||
Usage:
|
||||
python bench_gdn_decode.py # default sweep
|
||||
python bench_gdn_decode.py --mode bench # benchmark only
|
||||
python bench_gdn_decode.py --mode correctness # correctness only
|
||||
python bench_gdn_decode.py --preset qwen3.5-35b # Qwen3.5-35B-A3B config
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.srt.layers.attention.fla.fused_recurrent import (
|
||||
fused_recurrent_gated_delta_rule_packed_decode,
|
||||
)
|
||||
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
|
||||
fused_sigmoid_gating_delta_rule_update,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Input factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_inputs(
|
||||
B: int,
|
||||
H: int,
|
||||
HV: int,
|
||||
K: int,
|
||||
V: int,
|
||||
pool_size: int,
|
||||
device: str,
|
||||
dtype: torch.dtype,
|
||||
seed: int = 42,
|
||||
):
|
||||
"""Create all input tensors for a single benchmark / correctness run."""
|
||||
torch.manual_seed(seed)
|
||||
|
||||
qkv_dim = 2 * H * K + HV * V
|
||||
mixed_qkv = torch.randn(B, qkv_dim, device=device, dtype=dtype)
|
||||
a = torch.randn(B, HV, device=device, dtype=dtype)
|
||||
b = torch.randn(B, HV, device=device, dtype=dtype)
|
||||
A_log = torch.randn(HV, device=device, dtype=dtype)
|
||||
dt_bias = torch.randn(HV, device=device, dtype=dtype)
|
||||
|
||||
ssm_states = torch.randn(pool_size, HV, V, K, device=device, dtype=dtype) * 0.1
|
||||
cache_indices = torch.arange(B, device=device, dtype=torch.int32)
|
||||
|
||||
cu_seqlens = torch.arange(B + 1, device=device, dtype=torch.long)
|
||||
|
||||
return dict(
|
||||
B=B,
|
||||
H=H,
|
||||
HV=HV,
|
||||
K=K,
|
||||
V=V,
|
||||
qkv_dim=qkv_dim,
|
||||
pool_size=pool_size,
|
||||
mixed_qkv=mixed_qkv,
|
||||
a=a,
|
||||
b=b,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
ssm_states=ssm_states,
|
||||
cache_indices=cache_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Runner wrappers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_baseline(inp):
|
||||
"""Baseline path: split → view → fused_sigmoid_gating_delta_rule_update.
|
||||
|
||||
This mirrors the FULL original decode path in GDNAttnBackend.forward_decode,
|
||||
including the split, view, and kernel call.
|
||||
"""
|
||||
B, H, HV, K, V = inp["B"], inp["H"], inp["HV"], inp["K"], inp["V"]
|
||||
mixed_qkv = inp["mixed_qkv"]
|
||||
ssm_states = inp["ssm_states"].clone()
|
||||
|
||||
# Step 1: split (same as forward_decode)
|
||||
q_flat, k_flat, v_flat = torch.split(mixed_qkv, [H * K, H * K, HV * V], dim=-1)
|
||||
|
||||
# Step 2: view + reshape (same as forward_decode)
|
||||
q = q_flat.view(1, B, H, K)
|
||||
k = k_flat.view(1, B, H, K)
|
||||
v = v_flat.view(1, B, HV, V)
|
||||
|
||||
# Step 3: fused gating + recurrent update
|
||||
o = fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=ssm_states,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"],
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
)
|
||||
|
||||
return o, ssm_states
|
||||
|
||||
|
||||
def run_packed(inp):
|
||||
"""Packed path: single fused kernel directly on mixed_qkv."""
|
||||
B, HV, K, V = inp["B"], inp["HV"], inp["K"], inp["V"]
|
||||
ssm_states = inp["ssm_states"].clone()
|
||||
out = inp["mixed_qkv"].new_empty(B, 1, HV, V)
|
||||
|
||||
fused_recurrent_gated_delta_rule_packed_decode(
|
||||
mixed_qkv=inp["mixed_qkv"],
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
scale=inp["K"] ** -0.5,
|
||||
initial_state=ssm_states,
|
||||
out=out,
|
||||
ssm_state_indices=inp["cache_indices"],
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
# Convert [B, 1, HV, V] → [1, B, HV, V] to match baseline layout
|
||||
return out.transpose(0, 1), ssm_states
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Correctness check
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
|
||||
"""Run correctness check for a single config. Returns True if PASS."""
|
||||
tag = f"B={B:>4} H={H:>2} HV={HV:>2} K={K:>3} V={V:>3} pool={pool_size:>4}"
|
||||
|
||||
inp = make_inputs(B, H, HV, K, V, pool_size, device, dtype, seed=seed)
|
||||
|
||||
o_baseline, state_baseline = run_baseline(inp)
|
||||
o_packed, state_packed = run_packed(inp)
|
||||
|
||||
# Output comparison
|
||||
atol = 2e-2 if dtype != torch.float32 else 1e-4
|
||||
rtol = 1e-2 if dtype != torch.float32 else 1e-4
|
||||
|
||||
try:
|
||||
torch.testing.assert_close(o_packed, o_baseline, atol=atol, rtol=rtol)
|
||||
output_ok = True
|
||||
except AssertionError as e:
|
||||
output_ok = False
|
||||
out_diff = (o_packed - o_baseline).abs().max().item()
|
||||
|
||||
# State comparison (only for slots that were updated)
|
||||
indices = inp["cache_indices"]
|
||||
try:
|
||||
torch.testing.assert_close(
|
||||
state_packed[indices], state_baseline[indices], atol=atol, rtol=rtol
|
||||
)
|
||||
state_ok = True
|
||||
except AssertionError:
|
||||
state_ok = False
|
||||
st_diff = (state_packed[indices] - state_baseline[indices]).abs().max().item()
|
||||
|
||||
passed = output_ok and state_ok
|
||||
|
||||
if passed:
|
||||
print(f" [PASS] {tag}")
|
||||
else:
|
||||
details = []
|
||||
if not output_ok:
|
||||
details.append(f"output max_diff={out_diff:.6f}")
|
||||
if not state_ok:
|
||||
details.append(f"state max_diff={st_diff:.6f}")
|
||||
print(f" [FAIL] {tag} ({', '.join(details)})")
|
||||
|
||||
return passed
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
|
||||
"""Benchmark baseline vs packed for a single config."""
|
||||
inp = make_inputs(B, H, HV, K, V, pool_size, device, dtype)
|
||||
|
||||
# ── Baseline: full path including split + view ──
|
||||
def fn_baseline():
|
||||
q_flat, k_flat, v_flat = torch.split(
|
||||
inp["mixed_qkv"], [H * K, H * K, HV * V], dim=-1
|
||||
)
|
||||
q = q_flat.view(1, B, H, K)
|
||||
k = k_flat.view(1, B, H, K)
|
||||
v = v_flat.view(1, B, HV, V)
|
||||
fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
initial_state_source=inp["ssm_states"],
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"],
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
)
|
||||
|
||||
# ── Packed: single kernel ──
|
||||
out_buf = inp["mixed_qkv"].new_empty(B, 1, HV, V)
|
||||
|
||||
def fn_packed():
|
||||
fused_recurrent_gated_delta_rule_packed_decode(
|
||||
mixed_qkv=inp["mixed_qkv"],
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
A_log=inp["A_log"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
scale=K**-0.5,
|
||||
initial_state=inp["ssm_states"],
|
||||
out=out_buf,
|
||||
ssm_state_indices=inp["cache_indices"],
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
# Warmup
|
||||
for _ in range(10):
|
||||
fn_baseline()
|
||||
fn_packed()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
try:
|
||||
ms_baseline, ms_base_lo, ms_base_hi = triton.testing.do_bench(
|
||||
fn_baseline, quantiles=quantiles, warmup=50, rep=200
|
||||
)
|
||||
ms_packed, ms_pack_lo, ms_pack_hi = triton.testing.do_bench(
|
||||
fn_packed, quantiles=quantiles, warmup=50, rep=200
|
||||
)
|
||||
except Exception:
|
||||
# Fallback to manual timing
|
||||
torch.cuda.synchronize()
|
||||
N = 200
|
||||
start = time.perf_counter()
|
||||
for _ in range(N):
|
||||
fn_baseline()
|
||||
torch.cuda.synchronize()
|
||||
ms_baseline = (time.perf_counter() - start) / N * 1000
|
||||
|
||||
start = time.perf_counter()
|
||||
for _ in range(N):
|
||||
fn_packed()
|
||||
torch.cuda.synchronize()
|
||||
ms_packed = (time.perf_counter() - start) / N * 1000
|
||||
|
||||
speedup = ms_baseline / ms_packed if ms_packed > 0 else float("inf")
|
||||
saved_us = (ms_baseline - ms_packed) * 1000
|
||||
|
||||
print(
|
||||
f" {B:>5} {H:>3} {HV:>3} {K:>3} {V:>3} | "
|
||||
f"{ms_baseline * 1000:>10.1f} | "
|
||||
f"{ms_packed * 1000:>10.1f} | "
|
||||
f"{speedup:>7.2f}x | "
|
||||
f"{saved_us:>+9.1f}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_correctness(device, dtype):
|
||||
print("=" * 70)
|
||||
print("Correctness: Baseline GDN Decode vs Packed GDN Decode")
|
||||
print("=" * 70)
|
||||
|
||||
shapes = [
|
||||
# (B, H, HV, K, V, pool_size)
|
||||
# --- Qwen3.5-35B-A3B style (TP=2: H=8, HV=16) ---
|
||||
(1, 8, 16, 128, 128, 32),
|
||||
(4, 8, 16, 128, 128, 32),
|
||||
(16, 8, 16, 128, 128, 64),
|
||||
(32, 8, 16, 128, 128, 128),
|
||||
(64, 8, 16, 128, 128, 128),
|
||||
(128, 8, 16, 128, 128, 256),
|
||||
(256, 8, 16, 128, 128, 512),
|
||||
# --- Qwen3.5-35B-A3B style (TP=1: H=16, HV=32) ---
|
||||
(1, 16, 32, 128, 128, 32),
|
||||
(32, 16, 32, 128, 128, 128),
|
||||
(64, 16, 32, 128, 128, 128),
|
||||
# --- Qwen3-Next-80B-A3B style ---
|
||||
(32, 16, 16, 128, 128, 128),
|
||||
(64, 16, 16, 128, 128, 128),
|
||||
# --- With PAD_SLOT_ID ---
|
||||
(32, 8, 16, 128, 128, 128), # some indices may be padded
|
||||
# --- Edge cases ---
|
||||
(1, 8, 16, 128, 128, 32),
|
||||
(2, 8, 16, 128, 128, 32),
|
||||
]
|
||||
|
||||
all_pass = True
|
||||
for B, H, HV, K, V, pool_size in shapes:
|
||||
if not check_correctness(B, H, HV, K, V, pool_size, device, dtype):
|
||||
all_pass = False
|
||||
|
||||
# PAD_SLOT_ID test
|
||||
print("\n PAD_SLOT_ID test (indices with -1):")
|
||||
inp = make_inputs(32, 8, 16, 128, 128, 128, device, dtype)
|
||||
o_baseline, st_baseline = run_baseline(inp)
|
||||
o_packed, st_packed = run_packed(inp)
|
||||
|
||||
try:
|
||||
torch.testing.assert_close(o_packed, o_baseline, atol=2e-2, rtol=1e-2)
|
||||
print(" [PASS] PAD_SLOT_ID=-1 handling")
|
||||
except AssertionError:
|
||||
print(" [FAIL] PAD_SLOT_ID=-1 handling")
|
||||
all_pass = False
|
||||
|
||||
print()
|
||||
if all_pass:
|
||||
print("ALL PASSED.")
|
||||
else:
|
||||
print("SOME FAILED.")
|
||||
return all_pass
|
||||
|
||||
|
||||
def run_benchmark(device, dtype, args):
|
||||
print()
|
||||
print("=" * 85)
|
||||
print("Benchmark: Baseline GDN Decode vs Packed GDN Decode")
|
||||
print("=" * 85)
|
||||
|
||||
K = args.head_size_k
|
||||
V = args.head_size_v
|
||||
pool_size = args.pool_size
|
||||
|
||||
if args.preset == "qwen3.5-35b":
|
||||
# Qwen3.5-35B-A3B: H_qk=16, H_v=32, K=128, V=128
|
||||
# After TP=2: H=8, HV=16
|
||||
bench_configs = [
|
||||
# (B, H, HV) — TP=2 config
|
||||
(1, 8, 16),
|
||||
(2, 8, 16),
|
||||
(4, 8, 16),
|
||||
(8, 8, 16),
|
||||
(16, 8, 16),
|
||||
(32, 8, 16),
|
||||
(64, 8, 16),
|
||||
(128, 8, 16),
|
||||
(256, 8, 16),
|
||||
(512, 8, 16),
|
||||
# TP=1 config (full heads)
|
||||
(1, 16, 32),
|
||||
(8, 16, 32),
|
||||
(32, 16, 32),
|
||||
(64, 16, 32),
|
||||
(128, 16, 32),
|
||||
(256, 16, 32),
|
||||
]
|
||||
elif args.preset == "qwen3-next-80b":
|
||||
bench_configs = [
|
||||
# Qwen3-Next-80B-A3B: all same H=HV=16 after TP
|
||||
(1, 16, 16),
|
||||
(8, 16, 16),
|
||||
(32, 16, 16),
|
||||
(64, 16, 16),
|
||||
(128, 16, 16),
|
||||
(256, 16, 16),
|
||||
]
|
||||
else:
|
||||
bench_configs = []
|
||||
for B in args.batch_sizes:
|
||||
for H in args.num_q_heads:
|
||||
for HV in args.num_v_heads:
|
||||
bench_configs.append((B, H, HV))
|
||||
|
||||
print(f" Config: K={K}, V={V}, pool_size={pool_size}, dtype={dtype}")
|
||||
print(
|
||||
f" {'B':>5} {'H':>3} {'HV':>3} {'K':>3} {'V':>3} | "
|
||||
f"{'base (μs)':>10} | "
|
||||
f"{'packed (μs)':>10} | "
|
||||
f"{'speedup':>8} | "
|
||||
f"{'saved (μs)':>10}"
|
||||
)
|
||||
print(" " + "-" * 75)
|
||||
|
||||
for B, H, HV in bench_configs:
|
||||
actual_pool = max(pool_size, B + 16)
|
||||
bench_shape(B, H, HV, K, V, actual_pool, device, dtype)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark & Correctness: GDN Packed Decode vs Baseline"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=["all", "correctness", "bench"],
|
||||
default="all",
|
||||
help="Run mode (default: all)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset",
|
||||
choices=["qwen3.5-35b", "qwen3-next-80b", "custom"],
|
||||
default="qwen3.5-35b",
|
||||
help="Preset config (default: qwen3.5-35b)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
choices=["float16", "bfloat16", "float32"],
|
||||
default="bfloat16",
|
||||
)
|
||||
parser.add_argument("--head-size-k", type=int, default=128)
|
||||
parser.add_argument("--head-size-v", type=int, default=128)
|
||||
parser.add_argument("--pool-size", type=int, default=512)
|
||||
parser.add_argument(
|
||||
"--batch-sizes",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1, 4, 8, 16, 32, 64, 128, 256, 512],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-q-heads",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[8, 16],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-v-heads",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[16, 32],
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
device = "cuda"
|
||||
dtype = getattr(torch, args.dtype)
|
||||
|
||||
cap = torch.cuda.get_device_capability()
|
||||
dev_name = torch.cuda.get_device_name()
|
||||
print(f"Device: {dev_name} (SM {cap[0]}{cap[1]})")
|
||||
|
||||
if args.mode in ("all", "correctness"):
|
||||
all_pass = run_correctness(device, dtype)
|
||||
if not all_pass and args.mode == "all":
|
||||
print("\nSkipping benchmark due to correctness failures.")
|
||||
return 1
|
||||
|
||||
if args.mode in ("all", "bench"):
|
||||
run_benchmark(device, dtype, args)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
639
third_party/sglang/benchmark/bench_linear_attention/bench_gdn_prefill.py
vendored
Normal file
639
third_party/sglang/benchmark/bench_linear_attention/bench_gdn_prefill.py
vendored
Normal file
@@ -0,0 +1,639 @@
|
||||
"""
|
||||
Benchmark & Correctness: Triton GDN vs FlashInfer GDN (prefill).
|
||||
|
||||
Compares:
|
||||
- Triton: sglang's chunk_gated_delta_rule (K-contiguous pool, pool-indexed)
|
||||
- FlashInfer: flashinfer's chunk_gated_delta_rule (gather/scatter, 3D tensors)
|
||||
|
||||
The two kernels have different APIs:
|
||||
- Triton: q/k/v=[1,T,H,D], g=logsigmoid, beta=sigmoid, has initial_state_indices
|
||||
- FlashInfer: q/k/v=[T,H,D], g=alpha(float32), beta=float32, no indices (gathered state)
|
||||
|
||||
Reports correctness (output & state matching) and performance (ms, TFLOPS, TB/s).
|
||||
|
||||
Usage:
|
||||
python benchmark_gdn_prefill.py # default sweep
|
||||
python benchmark_gdn_prefill.py --mode bench # benchmark only
|
||||
python benchmark_gdn_prefill.py --mode correctness # correctness only
|
||||
python benchmark_gdn_prefill.py --preset qwen3-next # Qwen3-Next config
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))
|
||||
|
||||
import torch
|
||||
from flashinfer.gdn_prefill import (
|
||||
chunk_gated_delta_rule as flashinfer_chunk_gated_delta_rule,
|
||||
)
|
||||
|
||||
from sglang.srt.layers.attention.fla.chunk import (
|
||||
chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
|
||||
)
|
||||
from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_k_contiguous(t: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Given a V-contiguous tensor [..., K, V], return a K-contiguous view of the
|
||||
same logical shape [..., K, V] (physically [..., V, K], K-last).
|
||||
"""
|
||||
return t.transpose(-2, -1).contiguous().transpose(-2, -1)
|
||||
|
||||
|
||||
def gdn_flops(
|
||||
total_seq_len: int,
|
||||
num_heads: int,
|
||||
head_size_k: int,
|
||||
head_size_v: int,
|
||||
) -> int:
|
||||
"""
|
||||
FLOPs for GDN prefill (delta rule).
|
||||
|
||||
Per token per head:
|
||||
1. k @ v^T (outer product): 2 * K * V
|
||||
2. q @ state (output): 2 * K * V
|
||||
"""
|
||||
outer_product_flops = 2 * total_seq_len * num_heads * head_size_k * head_size_v
|
||||
output_flops = 2 * total_seq_len * num_heads * head_size_k * head_size_v
|
||||
return outer_product_flops + output_flops
|
||||
|
||||
|
||||
def gdn_bytes(
|
||||
total_seq_len: int,
|
||||
num_q_heads: int,
|
||||
num_v_heads: int,
|
||||
head_size_k: int,
|
||||
head_size_v: int,
|
||||
num_seqs: int,
|
||||
dtype: torch.dtype,
|
||||
) -> int:
|
||||
"""Memory bytes accessed (inputs + outputs + state)."""
|
||||
num_o_heads = max(num_q_heads, num_v_heads)
|
||||
elem = dtype.itemsize
|
||||
|
||||
q_bytes = total_seq_len * num_q_heads * head_size_k * elem
|
||||
k_bytes = total_seq_len * num_v_heads * head_size_k * elem
|
||||
v_bytes = total_seq_len * num_v_heads * head_size_v * elem
|
||||
o_bytes = total_seq_len * num_o_heads * head_size_v * elem
|
||||
|
||||
# state (float32): read + write
|
||||
state_bytes = 2 * num_seqs * num_o_heads * head_size_k * head_size_v * 4
|
||||
|
||||
# g, beta (float32)
|
||||
g_bytes = total_seq_len * num_o_heads * 4
|
||||
beta_bytes = total_seq_len * num_o_heads * 4
|
||||
|
||||
return q_bytes + k_bytes + v_bytes + o_bytes + state_bytes + g_bytes + beta_bytes
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Input factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_inputs(
|
||||
B: int,
|
||||
T_per_seq: int,
|
||||
H: int,
|
||||
K: int,
|
||||
V: int,
|
||||
pool_size: int,
|
||||
device: str,
|
||||
dtype: torch.dtype,
|
||||
sequential_indices: bool = False,
|
||||
seed: int = 42,
|
||||
):
|
||||
"""Create all input tensors for a single benchmark / correctness run.
|
||||
|
||||
Returns a dict with both Triton-format and FlashInfer-format tensors.
|
||||
"""
|
||||
T = B * T_per_seq
|
||||
torch.manual_seed(seed)
|
||||
|
||||
if sequential_indices:
|
||||
cache_indices = torch.arange(B, dtype=torch.int32, device=device)
|
||||
else:
|
||||
perm = torch.randperm(pool_size, device=device)[:B]
|
||||
cache_indices = perm.to(torch.int32)
|
||||
|
||||
pool_init = torch.randn(pool_size, H, K, V, dtype=dtype, device=device) * 0.1
|
||||
|
||||
cu_seqlens = torch.arange(
|
||||
0, (B + 1) * T_per_seq, T_per_seq, dtype=torch.long, device=device
|
||||
)
|
||||
|
||||
# Triton format: [1, T, H, D]
|
||||
q = torch.randn(1, T, H, K, dtype=dtype, device=device)
|
||||
k = torch.randn(1, T, H, K, dtype=dtype, device=device)
|
||||
v = torch.randn(1, T, H, V, dtype=dtype, device=device)
|
||||
|
||||
# g (logsigmoid) and beta (sigmoid) in Triton format: [1, T, H]
|
||||
g_raw = torch.randn(1, T, H, dtype=dtype, device=device)
|
||||
g_triton = torch.nn.functional.logsigmoid(g_raw) # logsigmoid for Triton
|
||||
beta_triton = torch.sigmoid(torch.randn(1, T, H, dtype=dtype, device=device))
|
||||
|
||||
return dict(
|
||||
B=B,
|
||||
T=T,
|
||||
T_per_seq=T_per_seq,
|
||||
H=H,
|
||||
K=K,
|
||||
V=V,
|
||||
pool_size=pool_size,
|
||||
cache_indices=cache_indices,
|
||||
pool_init=pool_init,
|
||||
cu_seqlens=cu_seqlens,
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g_triton=g_triton,
|
||||
beta_triton=beta_triton,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Runner wrappers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_triton(inp):
|
||||
"""Triton path: K-contiguous pool, pool-indexed, [1,T,H,D] tensors."""
|
||||
pool = make_k_contiguous(inp["pool_init"].clone())
|
||||
|
||||
o, _, h = triton_chunk_gated_delta_rule(
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=inp["v"],
|
||||
g=inp["g_triton"],
|
||||
beta=inp["beta_triton"],
|
||||
initial_state=pool,
|
||||
initial_state_indices=inp["cache_indices"],
|
||||
cu_seqlens=inp["cu_seqlens"],
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
return o, pool, h
|
||||
|
||||
|
||||
def run_flashinfer(inp):
|
||||
"""FlashInfer path: matches sglang FlashInferGDNKernel.extend() exactly.
|
||||
|
||||
Key differences from Triton path:
|
||||
- q, k are L2-normalized BEFORE calling the kernel
|
||||
- use_qk_l2norm_in_kernel=False (kernel skips internal normalization)
|
||||
- Tensors are [T, H, D] (no batch dim)
|
||||
- g is alpha = exp(logsigmoid(...)) = sigmoid(...), float32
|
||||
- beta is float32
|
||||
- initial_state is gathered from pool (no pool-index support)
|
||||
- Uses keyword arguments (matching sglang production code)
|
||||
|
||||
NOTE: FlashInfer GDN requires K == V (square head_size).
|
||||
"""
|
||||
K = inp["K"]
|
||||
V = inp["V"]
|
||||
assert K == V, f"FlashInfer GDN requires K == V, got K={K}, V={V}"
|
||||
|
||||
pool = make_k_contiguous(inp["pool_init"].clone())
|
||||
cache_indices = inp["cache_indices"]
|
||||
|
||||
# Gather states from K-contiguous pool -> K-contiguous float32
|
||||
# In production, ssm_states is already float32 so .float() is no-op.
|
||||
# Here pool_init is bf16, so .float() loses K-contiguous layout.
|
||||
gathered = pool[cache_indices]
|
||||
initial_state = make_k_contiguous(gathered.float().contiguous())
|
||||
|
||||
q_fi = l2norm_fwd(inp["q"][0].contiguous())
|
||||
k_fi = l2norm_fwd(inp["k"][0].contiguous())
|
||||
v_fi = inp["v"][0].contiguous()
|
||||
|
||||
# g -> alpha (exp of logsigmoid = sigmoid), float32
|
||||
alpha_fi = torch.exp(inp["g_triton"][0].to(torch.float32))
|
||||
# beta -> float32
|
||||
beta_fi = inp["beta_triton"][0].to(torch.float32)
|
||||
|
||||
cu_seqlens_fi = inp["cu_seqlens"].to(torch.int64)
|
||||
|
||||
# Call FlashInfer with keyword args (matching sglang production code)
|
||||
# use_qk_l2norm_in_kernel=False because we pre-normalized above
|
||||
o_fi, state_fi = flashinfer_chunk_gated_delta_rule(
|
||||
q=q_fi,
|
||||
k=k_fi,
|
||||
v=v_fi,
|
||||
g=alpha_fi,
|
||||
beta=beta_fi,
|
||||
scale=None,
|
||||
initial_state=initial_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=cu_seqlens_fi,
|
||||
use_qk_l2norm_in_kernel=False,
|
||||
)
|
||||
|
||||
# Scatter updated states back to K-contiguous pool
|
||||
pool[cache_indices] = state_fi.to(pool.dtype)
|
||||
|
||||
# Reshape output: [T, H, D] -> [1, T, H, D] to match Triton
|
||||
o_out = o_fi.unsqueeze(0)
|
||||
|
||||
return o_out, pool, state_fi
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Correctness check
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def check_shape(
|
||||
B,
|
||||
T_per_seq,
|
||||
H,
|
||||
K,
|
||||
V,
|
||||
pool_size,
|
||||
device,
|
||||
dtype,
|
||||
sequential_indices=False,
|
||||
seed=42,
|
||||
):
|
||||
"""Run correctness check for a single shape config. Returns True if PASS.
|
||||
|
||||
Pass/fail is based on OUTPUT comparison only (atol=5e-2).
|
||||
Pool state diff is reported as informational — state divergence over many
|
||||
tokens is expected due to different chunk sizes and accumulation order.
|
||||
"""
|
||||
tag = (
|
||||
f"B={B:>3} T/seq={T_per_seq:>4} H={H:>2} K={K:>3} V={V:>3} pool={pool_size:>4}"
|
||||
)
|
||||
idx_tag = " (seq)" if sequential_indices else ""
|
||||
|
||||
# FlashInfer GDN requires K == V (square head_size)
|
||||
if K != V:
|
||||
print(f" [SKIP] {tag}{idx_tag} (FlashInfer requires K==V)")
|
||||
return True
|
||||
|
||||
# FlashInfer GDN CUTLASS kernels are only compiled for head_size=128.
|
||||
# Running with other sizes causes illegal memory access that poisons
|
||||
# the CUDA context (unrecoverable), so we must skip upfront.
|
||||
FLASHINFER_SUPPORTED_HEAD_SIZES = {128}
|
||||
if K not in FLASHINFER_SUPPORTED_HEAD_SIZES:
|
||||
print(
|
||||
f" [SKIP] {tag}{idx_tag} (FlashInfer only supports head_size={FLASHINFER_SUPPORTED_HEAD_SIZES})"
|
||||
)
|
||||
return True
|
||||
|
||||
inp = make_inputs(
|
||||
B,
|
||||
T_per_seq,
|
||||
H,
|
||||
K,
|
||||
V,
|
||||
pool_size,
|
||||
device,
|
||||
dtype,
|
||||
sequential_indices=sequential_indices,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
o_triton, pool_triton, h_triton = run_triton(inp)
|
||||
|
||||
# FlashInfer may not support all head_size values (e.g., only 128).
|
||||
# CUDA errors from unsupported configs are often asynchronous, so we
|
||||
# must synchronize inside the try block to catch them here.
|
||||
try:
|
||||
o_fi, pool_fi, _ = run_flashinfer(inp)
|
||||
torch.cuda.synchronize()
|
||||
except Exception as e:
|
||||
# Catch RuntimeError, torch.AcceleratorError, etc.
|
||||
# Reset CUDA error state so subsequent tests can proceed
|
||||
try:
|
||||
torch.cuda.synchronize()
|
||||
except Exception:
|
||||
pass
|
||||
print(f" [SKIP] {tag}{idx_tag} (FlashInfer error: {e})")
|
||||
return True
|
||||
|
||||
cache_indices = inp["cache_indices"]
|
||||
|
||||
# --- Output comparison ---
|
||||
# bf16 prefill with L2norm + chunked accumulation
|
||||
torch.testing.assert_close(o_triton, o_fi, atol=5e-2, rtol=1e-2)
|
||||
|
||||
# --- Stride check ---
|
||||
def strides_ok(pool):
|
||||
s = pool.stride()
|
||||
return s[-2] == 1 and s[-1] == K
|
||||
|
||||
strides_triton = strides_ok(pool_triton)
|
||||
strides_fi = strides_ok(pool_fi)
|
||||
|
||||
passed = strides_triton and strides_fi
|
||||
|
||||
# Build detail string
|
||||
details = []
|
||||
if not strides_triton:
|
||||
details.append("triton strides bad")
|
||||
if not strides_fi:
|
||||
details.append("flashinfer strides bad")
|
||||
|
||||
status = "PASS" if passed else "FAIL"
|
||||
detail_str = f" [{', '.join(details)}]"
|
||||
print(f" [{status}] {tag}{idx_tag}")
|
||||
return passed
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bench_shape(B, H, T_per_seq, K, V, pool_size, device, dtype):
|
||||
"""Benchmark Triton vs FlashInfer for a single config. Requires K == V."""
|
||||
import triton.testing
|
||||
|
||||
assert K == V, f"FlashInfer GDN requires K == V, got K={K}, V={V}"
|
||||
|
||||
T = B * T_per_seq
|
||||
inp = make_inputs(B, T_per_seq, H, K, V, pool_size, device, dtype)
|
||||
|
||||
# -- Shared read-only tensors --
|
||||
q, k_t, v = inp["q"], inp["k"], inp["v"]
|
||||
g_triton, beta_triton = inp["g_triton"], inp["beta_triton"]
|
||||
cu_seqlens = inp["cu_seqlens"]
|
||||
cache_indices = inp["cache_indices"]
|
||||
seq_indices = torch.arange(B, dtype=torch.int32, device=device)
|
||||
pool_v = inp["pool_init"]
|
||||
|
||||
def fn_triton():
|
||||
pool = make_k_contiguous(pool_v.clone())
|
||||
triton_chunk_gated_delta_rule(
|
||||
q=q,
|
||||
k=k_t,
|
||||
v=v,
|
||||
g=g_triton,
|
||||
beta=beta_triton,
|
||||
initial_state=pool,
|
||||
initial_state_indices=cache_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
def fn_flashinfer():
|
||||
# -- Pre-compute FlashInfer format tensors (outside timing) --
|
||||
# Pre-normalize q and k (matching sglang production: l2norm_fwd)
|
||||
# q_fi = torch.nn.functional.normalize(q[0].contiguous().float(), p=2.0, dim=-1).to(
|
||||
# dtype
|
||||
# )
|
||||
# k_fi = torch.nn.functional.normalize(k_t[0].contiguous().float(), p=2.0, dim=-1).to(
|
||||
# dtype
|
||||
# )
|
||||
q_fi = l2norm_fwd(q[0].contiguous())
|
||||
k_fi = l2norm_fwd(k_t[0].contiguous())
|
||||
v_fi = v[0].contiguous()
|
||||
alpha_fi = torch.exp(g_triton[0].to(torch.float32))
|
||||
beta_fi = beta_triton[0].to(torch.float32)
|
||||
cu_seqlens_fi = cu_seqlens.to(torch.int64)
|
||||
pool = make_k_contiguous(pool_v.clone())
|
||||
gathered = pool[cache_indices]
|
||||
initial_state = make_k_contiguous(gathered.float().contiguous())
|
||||
flashinfer_chunk_gated_delta_rule(
|
||||
q=q_fi,
|
||||
k=k_fi,
|
||||
v=v_fi,
|
||||
g=alpha_fi,
|
||||
beta=beta_fi,
|
||||
scale=None,
|
||||
initial_state=initial_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=cu_seqlens_fi,
|
||||
use_qk_l2norm_in_kernel=False,
|
||||
)
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
# Warmup
|
||||
fn_triton()
|
||||
fn_flashinfer()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
ms_triton, _, _ = triton.testing.do_bench_cudagraph(fn_triton, quantiles=quantiles)
|
||||
ms_fi, _, _ = triton.testing.do_bench_cudagraph(fn_flashinfer, quantiles=quantiles)
|
||||
|
||||
# Metrics
|
||||
num_o_heads = H
|
||||
flops = gdn_flops(T, num_o_heads, K, V)
|
||||
mem_bytes = gdn_bytes(T, H, H, K, V, B, dtype)
|
||||
|
||||
tflops_triton = flops / ms_triton / 1e9
|
||||
tflops_fi = flops / ms_fi / 1e9
|
||||
tb_s_triton = mem_bytes / ms_triton / 1e9
|
||||
tb_s_fi = mem_bytes / ms_fi / 1e9
|
||||
|
||||
speedup = ms_triton / ms_fi if ms_fi > 0 else float("inf")
|
||||
|
||||
print(
|
||||
f" {B:>5} {H:>3} {T_per_seq:>6} {T:>7} | "
|
||||
f"{ms_triton:>8.3f} {tflops_triton:>7.2f} {tb_s_triton:>7.2f} | "
|
||||
f"{ms_fi:>8.3f} {tflops_fi:>7.2f} {tb_s_fi:>7.2f} | "
|
||||
f"{speedup:>7.2f}x"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_correctness(device, dtype):
|
||||
print("=" * 78)
|
||||
print("Correctness sweep: Triton vs FlashInfer")
|
||||
print("=" * 78)
|
||||
|
||||
shapes = [
|
||||
# (B, T_per_seq, H, K, V, pool_size)
|
||||
# --- baseline (Qwen3-Next style) ---
|
||||
(4, 64, 16, 128, 128, 32),
|
||||
(4, 256, 16, 128, 128, 32),
|
||||
# --- different batch sizes ---
|
||||
(1, 128, 16, 128, 128, 32),
|
||||
(8, 128, 16, 128, 128, 64),
|
||||
(16, 64, 16, 128, 128, 128),
|
||||
(32, 32, 16, 128, 128, 256),
|
||||
# --- different head counts ---
|
||||
(4, 128, 4, 128, 128, 32),
|
||||
(4, 128, 8, 128, 128, 32),
|
||||
(4, 128, 16, 64, 64, 32),
|
||||
(4, 128, 32, 128, 128, 32),
|
||||
(4, 128, 64, 128, 128, 32),
|
||||
# --- short sequences ---
|
||||
(4, 1, 16, 128, 128, 32),
|
||||
(4, 7, 16, 128, 128, 32),
|
||||
(4, 16, 16, 128, 128, 32),
|
||||
# --- large pool (sparse access) ---
|
||||
(4, 128, 16, 128, 128, 512),
|
||||
# --- combined stress ---
|
||||
(32, 128, 32, 128, 128, 256),
|
||||
]
|
||||
|
||||
shapes_seq = [
|
||||
(8, 128, 16, 128, 128, 8),
|
||||
(4, 128, 32, 128, 128, 4),
|
||||
(4, 128, 64, 128, 128, 4),
|
||||
(32, 128, 32, 128, 128, 32),
|
||||
]
|
||||
|
||||
all_pass = True
|
||||
for B, T_per_seq, H, K, V, pool_size in shapes:
|
||||
if not check_shape(B, T_per_seq, H, K, V, pool_size, device, dtype):
|
||||
all_pass = False
|
||||
|
||||
print()
|
||||
print("Sequential-index variants:")
|
||||
for B, T_per_seq, H, K, V, pool_size in shapes_seq:
|
||||
if not check_shape(
|
||||
B,
|
||||
T_per_seq,
|
||||
H,
|
||||
K,
|
||||
V,
|
||||
pool_size,
|
||||
device,
|
||||
dtype,
|
||||
sequential_indices=True,
|
||||
):
|
||||
all_pass = False
|
||||
|
||||
print()
|
||||
if all_pass:
|
||||
print("ALL PASSED.")
|
||||
else:
|
||||
print("SOME FAILED.")
|
||||
return all_pass
|
||||
|
||||
|
||||
def run_benchmark(device, dtype, args):
|
||||
print()
|
||||
print("=" * 105)
|
||||
print("Benchmark: Triton GDN vs FlashInfer GDN (do_bench_cudagraph)")
|
||||
print("=" * 105)
|
||||
|
||||
K = args.head_size_k
|
||||
V = args.head_size_v
|
||||
pool_size = args.pool_size
|
||||
|
||||
if args.preset == "qwen3-next":
|
||||
bench_configs = [
|
||||
# (B, H, T_per_seq)
|
||||
(4, 16, 256),
|
||||
(4, 32, 256),
|
||||
(16, 16, 256),
|
||||
(16, 32, 256),
|
||||
(32, 16, 256),
|
||||
(32, 32, 256),
|
||||
(64, 16, 256),
|
||||
(64, 32, 256),
|
||||
(128, 16, 256),
|
||||
(128, 32, 256),
|
||||
# longer sequences
|
||||
(4, 16, 1024),
|
||||
(4, 32, 1024),
|
||||
(32, 16, 1024),
|
||||
(32, 32, 1024),
|
||||
]
|
||||
else:
|
||||
bench_configs = []
|
||||
for B in args.batch_sizes:
|
||||
for H in args.num_heads:
|
||||
for T_per_seq in args.seq_lens:
|
||||
bench_configs.append((B, H, T_per_seq))
|
||||
|
||||
print(f" Config: K={K}, V={V}, pool_size={pool_size}, dtype={dtype}")
|
||||
print(
|
||||
f" {'B':>5} {'H':>3} {'T/seq':>6} {'T_tot':>7} | "
|
||||
f"{'tri(ms)':>8} {'TFLOPS':>7} {'TB/s':>7} | "
|
||||
f"{'fi(ms)':>8} {'TFLOPS':>7} {'TB/s':>7} | "
|
||||
f"{'speedup':>8}"
|
||||
)
|
||||
print(" " + "-" * 98)
|
||||
|
||||
for B, H, T_per_seq in bench_configs:
|
||||
actual_pool = max(pool_size, B)
|
||||
bench_shape(B, H, T_per_seq, K, V, actual_pool, device, dtype)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark & Correctness: Triton GDN vs FlashInfer GDN"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=["all", "correctness", "bench"],
|
||||
default="all",
|
||||
help="Run mode (default: all)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset",
|
||||
choices=["qwen3-next", "custom"],
|
||||
default="qwen3-next",
|
||||
help="Preset config (default: qwen3-next)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
choices=["float16", "bfloat16"],
|
||||
default="bfloat16",
|
||||
)
|
||||
parser.add_argument("--head-size-k", type=int, default=128)
|
||||
parser.add_argument("--head-size-v", type=int, default=128)
|
||||
parser.add_argument("--pool-size", type=int, default=256)
|
||||
parser.add_argument(
|
||||
"--batch-sizes",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[4, 16, 32, 64, 128],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-heads",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[16, 32],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seq-lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[128, 256, 512, 1024],
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.preset == "qwen3-next":
|
||||
args.head_size_k = 128
|
||||
args.head_size_v = 128
|
||||
|
||||
device = "cuda"
|
||||
dtype = getattr(torch, args.dtype)
|
||||
|
||||
# Check SM version
|
||||
cap = torch.cuda.get_device_capability()
|
||||
dev_name = torch.cuda.get_device_name()
|
||||
print(f"Device: {dev_name} (SM {cap[0]}{cap[1]})")
|
||||
|
||||
if args.mode in ("all", "correctness"):
|
||||
all_pass = run_correctness(device, dtype)
|
||||
if not all_pass and args.mode == "all":
|
||||
print("\nSkipping benchmark due to correctness failures.")
|
||||
return 1
|
||||
|
||||
if args.mode in ("all", "bench"):
|
||||
run_benchmark(device, dtype, args)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user