mirror of
https://github.com/kvcache-ai/sglang.git
synced 2026-07-08 00:07:10 +00:00
300 lines
10 KiB
Python
300 lines
10 KiB
Python
"""Tests for CuTe DSL fused sigmoid gating delta rule kernel (GDN)."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
try:
|
|
import cuda.bindings.driver as cuda_driver
|
|
import cutlass # noqa: F401
|
|
from cutlass.cute.runtime import from_dlpack
|
|
|
|
from sglang.jit_kernel import cutedsl_gdn
|
|
|
|
CUTEDSL_AVAILABLE = True
|
|
except ImportError:
|
|
CUTEDSL_AVAILABLE = False
|
|
cutedsl_gdn = None
|
|
|
|
try:
|
|
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
|
|
fused_sigmoid_gating_delta_rule_update,
|
|
)
|
|
|
|
TRITON_AVAILABLE = True
|
|
except ImportError:
|
|
TRITON_AVAILABLE = False
|
|
|
|
|
|
def run_triton_kernel(A_log, dt_bias, q, k, v, a, b, initial_state, indices, scale):
|
|
return 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=initial_state,
|
|
initial_state_indices=indices,
|
|
scale=scale,
|
|
use_qk_l2norm_in_kernel=True,
|
|
cu_seqlens=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not CUTEDSL_AVAILABLE, reason="CuTe DSL not available")
|
|
@pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton kernel not available")
|
|
@pytest.mark.parametrize("B", [16, 128])
|
|
def test_cutedsl_gdn_precision(B: int):
|
|
"""Test precision of CuTe DSL GDN kernel against Triton reference."""
|
|
torch.manual_seed(2025)
|
|
T, H, K, V, HV = 1, 16, 128, 128, 32
|
|
scale = K**-0.5
|
|
|
|
A_log = torch.randn(HV, dtype=torch.float32, device="cuda")
|
|
dt_bias = torch.randn(HV, dtype=torch.bfloat16, device="cuda")
|
|
a = torch.randn(B, T, HV, dtype=torch.bfloat16, device="cuda")
|
|
b = torch.randn(B, T, HV, dtype=torch.bfloat16, device="cuda")
|
|
q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device="cuda")
|
|
k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device="cuda")
|
|
v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device="cuda")
|
|
indices = torch.arange(B, dtype=torch.int32, device="cuda")
|
|
state_cutedsl = torch.randn(B, HV, K, V, dtype=torch.float32, device="cuda")
|
|
state_triton = state_cutedsl.clone().reshape(-1).contiguous()
|
|
|
|
# Warmup compilation
|
|
_ = cutedsl_gdn.cutedsl_fused_sigmoid_gating_delta_rule_update(
|
|
A_log, dt_bias, q, k, v, a, b, state_cutedsl.clone(), indices, scale=scale
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# Fresh state for actual test
|
|
state_cutedsl = torch.randn(B, HV, K, V, dtype=torch.float32, device="cuda")
|
|
state_triton = state_cutedsl.clone().reshape(-1).contiguous()
|
|
|
|
out_cutedsl = cutedsl_gdn.cutedsl_fused_sigmoid_gating_delta_rule_update(
|
|
A_log, dt_bias, q, k, v, a, b, state_cutedsl, indices, scale=scale
|
|
)
|
|
out_triton = run_triton_kernel(
|
|
A_log, dt_bias, q, k, v, a, b, state_triton, indices, scale
|
|
)
|
|
|
|
# Check precision: diff > 0.1 must be < 1% of elements
|
|
abs_diff = (out_triton.float() - out_cutedsl.float()).abs()
|
|
max_diff = abs_diff.max().item()
|
|
mean_diff = abs_diff.mean().item()
|
|
fail_rate = (abs_diff > 0.1).float().mean().item() * 100
|
|
has_nan = torch.isnan(out_cutedsl).any() or torch.isinf(out_cutedsl).any()
|
|
|
|
kernel_type = "SmallBatch" if B < 32 else "LargeBatch"
|
|
print(
|
|
f"\n B={B} ({kernel_type}): max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e}, fail_rate={fail_rate:.2f}%"
|
|
)
|
|
|
|
assert not has_nan, "Output contains NaN/Inf"
|
|
assert fail_rate < 1.0, f"Fail rate {fail_rate:.2f}% >= 1%"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
True,
|
|
reason="Skip the performance test because the speedup ratio is highly unstable in the CI environment. ",
|
|
)
|
|
@pytest.mark.skipif(not CUTEDSL_AVAILABLE, reason="CuTe DSL not available")
|
|
@pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton kernel not available")
|
|
@pytest.mark.parametrize("B", [1, 128])
|
|
def test_cutedsl_gdn_performance(B: int):
|
|
"""Benchmark CuTe DSL GDN kernel against Triton reference."""
|
|
torch.manual_seed(2025)
|
|
T, H, K, V, HV = 1, 16, 128, 128, 32
|
|
N = B
|
|
scale = K**-0.5
|
|
is_varlen = True
|
|
warmup, bench_iters, run_iters = 10, 100, 10
|
|
|
|
A_log = torch.randn(HV, dtype=torch.float32, device="cuda")
|
|
dt_bias = torch.randn(HV, dtype=torch.bfloat16, device="cuda")
|
|
indices = torch.arange(N, dtype=torch.int32, device="cuda")
|
|
state_cutedsl = torch.randn(N, HV, K, V, dtype=torch.float32, device="cuda")
|
|
state_triton = state_cutedsl.reshape(-1).contiguous()
|
|
cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda")
|
|
o_cutedsl = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda")
|
|
|
|
# Prepare tensors for multiple runs
|
|
q_list, k_list, v_list, a_list, b_list = [], [], [], [], []
|
|
q_tensor_list, k_tensor_list, v_tensor_list, a_tensor_list, b_tensor_list = (
|
|
[],
|
|
[],
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
q_triton, k_triton, v_triton, a_triton, b_triton = [], [], [], [], []
|
|
|
|
for ri in range(run_iters):
|
|
torch.manual_seed(2025 + ri)
|
|
q_i = torch.randn(1, N, H, K, dtype=torch.bfloat16, device="cuda")
|
|
k_i = torch.randn(1, N, H, K, dtype=torch.bfloat16, device="cuda")
|
|
v_i = torch.randn(1, N, HV, V, dtype=torch.bfloat16, device="cuda")
|
|
a_i = torch.randn(N, HV, dtype=torch.bfloat16, device="cuda")
|
|
b_i = torch.randn(N, HV, dtype=torch.bfloat16, device="cuda")
|
|
|
|
q_list.append(q_i)
|
|
k_list.append(k_i)
|
|
v_list.append(v_i)
|
|
a_list.append(a_i)
|
|
b_list.append(b_i)
|
|
q_tensor_list.append(from_dlpack(q_i, assumed_align=16))
|
|
k_tensor_list.append(from_dlpack(k_i, assumed_align=16))
|
|
v_tensor_list.append(from_dlpack(v_i, assumed_align=16))
|
|
a_tensor_list.append(from_dlpack(a_i, assumed_align=16))
|
|
b_tensor_list.append(from_dlpack(b_i, assumed_align=16))
|
|
q_triton.append(q_i.transpose(0, 1).contiguous())
|
|
k_triton.append(k_i.transpose(0, 1).contiguous())
|
|
v_triton.append(v_i.transpose(0, 1).contiguous())
|
|
a_triton.append(a_i.unsqueeze(1).contiguous())
|
|
b_triton.append(b_i.unsqueeze(1).contiguous())
|
|
|
|
A_log_t = from_dlpack(A_log, assumed_align=16)
|
|
dt_bias_t = from_dlpack(dt_bias, assumed_align=16)
|
|
h0_t = from_dlpack(state_cutedsl, assumed_align=16)
|
|
idx_t = from_dlpack(indices, assumed_align=16)
|
|
o_t = from_dlpack(o_cutedsl, assumed_align=16)
|
|
cu_t = from_dlpack(cu_seqlens, assumed_align=16)
|
|
|
|
torch_stream = torch.cuda.Stream()
|
|
stream = cuda_driver.CUstream(torch_stream.cuda_stream)
|
|
|
|
# Compile kernels
|
|
compiled = cutedsl_gdn._get_compiled_kernel(N, H, HV, K, V, N, N < 32, is_varlen)
|
|
torch.cuda.synchronize()
|
|
|
|
for ri in range(run_iters):
|
|
_ = run_triton_kernel(
|
|
A_log,
|
|
dt_bias,
|
|
q_triton[ri],
|
|
k_triton[ri],
|
|
v_triton[ri],
|
|
a_triton[ri],
|
|
b_triton[ri],
|
|
state_triton,
|
|
indices,
|
|
scale,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
def run_cutedsl():
|
|
for ri in range(run_iters):
|
|
compiled(
|
|
cu_t,
|
|
q_tensor_list[ri],
|
|
k_tensor_list[ri],
|
|
v_tensor_list[ri],
|
|
a_tensor_list[ri],
|
|
b_tensor_list[ri],
|
|
A_log_t,
|
|
dt_bias_t,
|
|
h0_t,
|
|
idx_t,
|
|
o_t,
|
|
stream,
|
|
)
|
|
|
|
def run_triton():
|
|
for ri in range(run_iters):
|
|
_ = run_triton_kernel(
|
|
A_log,
|
|
dt_bias,
|
|
q_triton[ri],
|
|
k_triton[ri],
|
|
v_triton[ri],
|
|
a_triton[ri],
|
|
b_triton[ri],
|
|
state_triton,
|
|
indices,
|
|
scale,
|
|
)
|
|
|
|
# Warmup
|
|
with torch.cuda.stream(torch_stream):
|
|
run_cutedsl()
|
|
torch.cuda.synchronize()
|
|
run_triton()
|
|
torch.cuda.synchronize()
|
|
|
|
# Capture CUDA graphs
|
|
graph_triton = torch.cuda.CUDAGraph()
|
|
graph_cutedsl = torch.cuda.CUDAGraph()
|
|
try:
|
|
with torch.cuda.graph(graph_triton):
|
|
run_triton()
|
|
with torch.cuda.graph(graph_cutedsl, stream=torch_stream):
|
|
run_cutedsl()
|
|
torch.cuda.synchronize()
|
|
except Exception:
|
|
graph_triton = graph_cutedsl = None
|
|
|
|
# Warmup with graphs
|
|
for _ in range(warmup):
|
|
if graph_cutedsl:
|
|
graph_cutedsl.replay()
|
|
else:
|
|
with torch.cuda.stream(torch_stream):
|
|
run_cutedsl()
|
|
torch.cuda.synchronize()
|
|
|
|
if graph_triton:
|
|
graph_triton.replay()
|
|
else:
|
|
run_triton()
|
|
torch.cuda.synchronize()
|
|
|
|
# Benchmark
|
|
triton_times, cutedsl_times = [], []
|
|
for _ in range(bench_iters):
|
|
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(
|
|
enable_timing=True
|
|
)
|
|
start.record()
|
|
if graph_triton:
|
|
graph_triton.replay()
|
|
else:
|
|
run_triton()
|
|
end.record()
|
|
torch.cuda.synchronize()
|
|
triton_times.append(start.elapsed_time(end))
|
|
|
|
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(
|
|
enable_timing=True
|
|
)
|
|
with torch.cuda.stream(torch_stream):
|
|
start.record()
|
|
if graph_cutedsl:
|
|
graph_cutedsl.replay()
|
|
else:
|
|
run_cutedsl()
|
|
end.record()
|
|
torch.cuda.synchronize()
|
|
cutedsl_times.append(start.elapsed_time(end))
|
|
|
|
triton_mean = np.mean(triton_times) / run_iters * 1000
|
|
triton_std = np.std(triton_times) / run_iters * 1000
|
|
cutedsl_mean = np.mean(cutedsl_times) / run_iters * 1000
|
|
cutedsl_std = np.std(cutedsl_times) / run_iters * 1000
|
|
speedup = triton_mean / cutedsl_mean
|
|
|
|
kernel_type = "SmallBatch" if B < 32 else "LargeBatch"
|
|
print(
|
|
f"\n B={B} ({kernel_type}): Triton={triton_mean:.2f}±{triton_std:.2f}μs, CuTeDSL={cutedsl_mean:.2f}±{cutedsl_std:.2f}μs, speedup={speedup:.2f}x"
|
|
)
|
|
|
|
min_speedup = 1.0 if B < 32 else 1.15
|
|
assert speedup >= min_speedup, f"Speedup {speedup:.2f}x < {min_speedup}x for B={B}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|