316 lines
9.8 KiB
Python
316 lines
9.8 KiB
Python
from typing import Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
# Import the function to benchmark
|
|
from sglang.srt.layers.attention.fla.layernorm_gated import (
|
|
_layer_norm_fwd as layer_norm_fwd,
|
|
)
|
|
from sglang.srt.layers.attention.fla.layernorm_gated import (
|
|
rms_norm_ref,
|
|
)
|
|
|
|
|
|
def benchmark_layer_norm_fwd(
|
|
M: int = 65536,
|
|
N: int = 128,
|
|
eps: float = 1e-6,
|
|
has_z: bool = True,
|
|
has_bias: bool = False,
|
|
group_size: Optional[int] = None,
|
|
norm_before_gate: bool = True,
|
|
is_rms_norm: bool = True,
|
|
dtype: torch.dtype = torch.float16,
|
|
warmup_iters: int = 10,
|
|
benchmark_iters: int = 100,
|
|
device: str = "cuda",
|
|
verbose: bool = True,
|
|
):
|
|
"""
|
|
Benchmark layer_norm_fwd with specified parameters.
|
|
|
|
Args:
|
|
M: Number of rows (batch size)
|
|
N: Number of columns (hidden dimension)
|
|
eps: Epsilon for numerical stability
|
|
has_z: Whether to use gating tensor z
|
|
has_bias: Whether to use bias
|
|
group_size: Group size for group normalization (None = full dimension)
|
|
norm_before_gate: Whether to normalize before gating
|
|
is_rms_norm: Whether to use RMS normalization (vs LayerNorm)
|
|
dtype: Data type for tensors
|
|
warmup_iters: Number of warmup iterations
|
|
benchmark_iters: Number of benchmark iterations
|
|
device: Device to run on
|
|
"""
|
|
if verbose:
|
|
print("=" * 80)
|
|
print("LayerNorm Forward Pass Benchmark")
|
|
print("=" * 80)
|
|
print(f"\nConfiguration:")
|
|
print(f" x.shape: torch.Size([{M}, {N}])")
|
|
print(f" weight.shape: torch.Size([{N}])")
|
|
print(f" bias: {'torch.Size([{}])'.format(N) if has_bias else None}")
|
|
print(f" eps: {eps}")
|
|
print(f" z: {'torch.Size([{}, {}])'.format(M, N) if has_z else None}")
|
|
print(f" out: None")
|
|
print(f" group_size: {group_size}")
|
|
print(f" norm_before_gate: {norm_before_gate}")
|
|
print(f" is_rms_norm: {is_rms_norm}")
|
|
print(f" dtype: {dtype}")
|
|
print(f" device: {device}")
|
|
print()
|
|
|
|
# Create input tensors
|
|
torch.manual_seed(42)
|
|
x = torch.randn(M, N, dtype=dtype, device=device)
|
|
weight = torch.randn(N, dtype=dtype, device=device)
|
|
bias = torch.randn(N, dtype=dtype, device=device) if has_bias else None
|
|
z = torch.randn(M, N, dtype=dtype, device=device) if has_z else None
|
|
|
|
# Ensure contiguous memory layout
|
|
x = x.contiguous()
|
|
weight = weight.contiguous()
|
|
if bias is not None:
|
|
bias = bias.contiguous()
|
|
if z is not None:
|
|
z = z.contiguous()
|
|
|
|
if verbose:
|
|
print("Warming up...")
|
|
# Warmup
|
|
for _ in range(warmup_iters):
|
|
out, mean, rstd = layer_norm_fwd(
|
|
x=x,
|
|
weight=weight,
|
|
bias=bias,
|
|
eps=eps,
|
|
z=z,
|
|
out=None,
|
|
group_size=group_size,
|
|
norm_before_gate=norm_before_gate,
|
|
is_rms_norm=is_rms_norm,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
if verbose:
|
|
print(f"Capturing CUDA graph...")
|
|
|
|
# Capture the kernel execution in a CUDA graph
|
|
runs_per_measurement = 100
|
|
|
|
# Create output tensor for graph capture
|
|
out_graph = torch.empty_like(x)
|
|
mean_graph = (
|
|
torch.empty((x.shape[0],), dtype=torch.float32, device=x.device)
|
|
if not is_rms_norm
|
|
else None
|
|
)
|
|
rstd_graph = torch.empty((x.shape[0],), dtype=torch.float32, device=x.device)
|
|
|
|
# Capture the graph
|
|
graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(graph):
|
|
for _ in range(runs_per_measurement):
|
|
out, mean, rstd = layer_norm_fwd(
|
|
x=x,
|
|
weight=weight,
|
|
bias=bias,
|
|
eps=eps,
|
|
z=z,
|
|
out=out_graph,
|
|
group_size=group_size,
|
|
norm_before_gate=norm_before_gate,
|
|
is_rms_norm=is_rms_norm,
|
|
)
|
|
|
|
if verbose:
|
|
print(
|
|
f"Running benchmark with {benchmark_iters} iterations using CUDA graph..."
|
|
)
|
|
|
|
# Benchmark by replaying the graph
|
|
times = []
|
|
for i in range(benchmark_iters):
|
|
start_event = torch.cuda.Event(enable_timing=True)
|
|
end_event = torch.cuda.Event(enable_timing=True)
|
|
|
|
start_event.record()
|
|
graph.replay()
|
|
end_event.record()
|
|
torch.cuda.synchronize()
|
|
|
|
# elapsed_time_ms returns milliseconds, divide by runs_per_measurement
|
|
elapsed_ms = start_event.elapsed_time(end_event)
|
|
times.append(
|
|
elapsed_ms / 1000.0 / runs_per_measurement
|
|
) # Convert to seconds per run
|
|
|
|
# Compute statistics
|
|
times = np.array(times) * 1_000_000 # Convert to microseconds
|
|
mean_time = np.mean(times)
|
|
std_time = np.std(times)
|
|
min_time = np.min(times)
|
|
max_time = np.max(times)
|
|
median_time = np.median(times)
|
|
p95_time = np.percentile(times, 95)
|
|
p99_time = np.percentile(times, 99)
|
|
|
|
# Calculate throughput
|
|
num_elements = M * N
|
|
throughput_gelements_per_sec = (num_elements / mean_time) * 1_000_000 / 1e9
|
|
|
|
# Calculate memory bandwidth
|
|
# Read: x, weight, z (if has_z)
|
|
# Write: out, rstd, mean (if not rms_norm)
|
|
bytes_per_element = 2 if dtype == torch.float16 else 4 # fp16 or fp32
|
|
read_bytes = (M * N + N) * bytes_per_element # x + weight
|
|
if has_z:
|
|
read_bytes += M * N * bytes_per_element # z
|
|
write_bytes = M * N * bytes_per_element # out
|
|
write_bytes += M * 4 # rstd (float32)
|
|
if not is_rms_norm:
|
|
write_bytes += M * 4 # mean (float32)
|
|
|
|
total_bytes = read_bytes + write_bytes
|
|
bandwidth_gb_per_sec = (total_bytes / mean_time) * 1_000_000 / 1e9
|
|
|
|
if verbose:
|
|
print("\n" + "=" * 80)
|
|
print("Benchmark Results")
|
|
print("=" * 80)
|
|
print(f"\nTiming Statistics (microseconds):")
|
|
print(f" Mean: {mean_time:.2f} us")
|
|
print(f" Std Dev: {std_time:.2f} us")
|
|
print(f" Min: {min_time:.2f} us")
|
|
print(f" Max: {max_time:.2f} us")
|
|
print(f" Median: {median_time:.2f} us")
|
|
print(f" P95: {p95_time:.2f} us")
|
|
print(f" P99: {p99_time:.2f} us")
|
|
|
|
print(f"\nThroughput:")
|
|
print(f" {throughput_gelements_per_sec:.2f} GElements/sec")
|
|
print(f" {bandwidth_gb_per_sec:.2f} GB/sec")
|
|
|
|
print(f"\nMemory Usage:")
|
|
print(f" Input size: {read_bytes / 1e6:.2f} MB")
|
|
print(f" Output size: {write_bytes / 1e6:.2f} MB")
|
|
print(f" Total: {total_bytes / 1e6:.2f} MB")
|
|
|
|
# Verify correctness against reference implementation
|
|
if verbose:
|
|
print("\nVerifying correctness...")
|
|
out_triton, mean_triton, rstd_triton = layer_norm_fwd(
|
|
x=x,
|
|
weight=weight,
|
|
bias=bias,
|
|
eps=eps,
|
|
z=z,
|
|
out=None,
|
|
group_size=group_size,
|
|
norm_before_gate=norm_before_gate,
|
|
is_rms_norm=is_rms_norm,
|
|
)
|
|
|
|
# Compute reference output
|
|
out_ref = rms_norm_ref(
|
|
x=x,
|
|
weight=weight,
|
|
bias=bias,
|
|
z=z,
|
|
eps=eps,
|
|
group_size=group_size,
|
|
norm_before_gate=norm_before_gate,
|
|
upcast=True,
|
|
)
|
|
|
|
# Compare outputs
|
|
max_diff = torch.max(torch.abs(out_triton - out_ref)).item()
|
|
mean_diff = torch.mean(torch.abs(out_triton - out_ref)).item()
|
|
rel_diff = torch.mean(
|
|
torch.abs(out_triton - out_ref) / (torch.abs(out_ref) + 1e-5)
|
|
).item()
|
|
|
|
if verbose:
|
|
print(f"\nCorrectness Check (vs Reference Implementation):")
|
|
print(f" Max absolute difference: {max_diff:.6e}")
|
|
print(f" Mean absolute difference: {mean_diff:.6e}")
|
|
print(f" Mean relative difference: {rel_diff:.6e}")
|
|
|
|
if max_diff < 1e-2:
|
|
print(" ✓ PASS: Results match reference implementation")
|
|
else:
|
|
print(" ✗ FAIL: Results do not match reference implementation")
|
|
|
|
print("\n" + "=" * 80)
|
|
|
|
return {
|
|
"mean_time_us": mean_time,
|
|
"std_time_us": std_time,
|
|
"min_time_us": min_time,
|
|
"max_time_us": max_time,
|
|
"median_time_us": median_time,
|
|
"p95_time_us": p95_time,
|
|
"p99_time_us": p99_time,
|
|
"throughput_gelements_per_sec": throughput_gelements_per_sec,
|
|
"bandwidth_gb_per_sec": bandwidth_gb_per_sec,
|
|
"max_diff": max_diff,
|
|
"mean_diff": mean_diff,
|
|
"rel_diff": rel_diff,
|
|
}
|
|
|
|
|
|
def main():
|
|
"""Run the benchmark with the specified configuration."""
|
|
# Configuration from user
|
|
config = {
|
|
"M": 65536,
|
|
"N": 128,
|
|
"eps": 1e-6,
|
|
"has_z": True,
|
|
"has_bias": False,
|
|
"group_size": None,
|
|
"norm_before_gate": True,
|
|
"is_rms_norm": True,
|
|
"dtype": torch.float16,
|
|
"warmup_iters": 10,
|
|
"benchmark_iters": 100,
|
|
"device": "cuda",
|
|
}
|
|
|
|
if not torch.cuda.is_available():
|
|
print("CUDA is not available. This benchmark requires a CUDA-enabled GPU.")
|
|
return
|
|
|
|
results = benchmark_layer_norm_fwd(**config)
|
|
|
|
# Collect all results
|
|
all_results = []
|
|
# Test with different batch sizes
|
|
print("\nRunning benchmarks for varying batch sizes...")
|
|
for M in [256, 512, 1024, 4096, 16384, 65536, 2**17, 2**18]:
|
|
config_var = config.copy()
|
|
config_var["M"] = M
|
|
config_var["warmup_iters"] = 5
|
|
config_var["benchmark_iters"] = 50
|
|
config_var["verbose"] = False
|
|
result = benchmark_layer_norm_fwd(**config_var)
|
|
all_results.append({"M": M, "N": config_var["N"], **result})
|
|
print(f" M={M:>5}: {result['mean_time_us']:>7.2f} us")
|
|
|
|
# Print summary table
|
|
print("\n\n")
|
|
print("=" * 30)
|
|
print("SUMMARY TABLE - Varying Batch Size (M) with N=128")
|
|
print("=" * 30)
|
|
print(f"{'M':>8} | {'Median (us)':>12}")
|
|
print("-" * 30)
|
|
for r in all_results:
|
|
print(f"{r['M']:>8} | {r['median_time_us']:>12.2f}")
|
|
print("=" * 30)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|