Optimize layernorm_gated for Qwen3-Next (#16397)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
313
benchmark/fla/benchmark_layernorm_gated.py
Normal file
313
benchmark/fla/benchmark_layernorm_gated.py
Normal file
@@ -0,0 +1,313 @@
|
||||
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()
|
||||
@@ -6,13 +6,15 @@
|
||||
# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.
|
||||
|
||||
|
||||
from functools import lru_cache
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from einops import rearrange
|
||||
|
||||
from sglang.srt.utils import device_context, is_npu
|
||||
from sglang.srt.utils import cdiv, device_context, is_npu, next_power_of_2
|
||||
|
||||
_is_npu = is_npu()
|
||||
|
||||
@@ -63,55 +65,103 @@ def _layer_norm_fwd_1pass_kernel(
|
||||
stride_y_row,
|
||||
stride_z_row,
|
||||
M, # number of rows in X
|
||||
N, # number of columns in X
|
||||
N: tl.constexpr, # number of columns in X
|
||||
eps, # epsilon to avoid division by zero
|
||||
BLOCK_N: tl.constexpr,
|
||||
ROWS_PER_BLOCK: tl.constexpr,
|
||||
HAS_BIAS: tl.constexpr,
|
||||
HAS_Z: tl.constexpr,
|
||||
NORM_BEFORE_GATE: tl.constexpr,
|
||||
IS_RMS_NORM: tl.constexpr,
|
||||
):
|
||||
# Map the program id to the row of X and Y it should compute.
|
||||
row = tl.program_id(0)
|
||||
# Map the program id to the starting row of X and Y it should compute.
|
||||
row_start = tl.program_id(0) * ROWS_PER_BLOCK
|
||||
group = tl.program_id(1)
|
||||
X += row * stride_x_row + group * N
|
||||
Y += row * stride_y_row + group * N
|
||||
if HAS_Z:
|
||||
Z += row * stride_z_row + group * N
|
||||
if not IS_RMS_NORM:
|
||||
Mean += group * M
|
||||
Rstd += group * M
|
||||
W += group * N
|
||||
if HAS_BIAS:
|
||||
B += group * N
|
||||
# Compute mean and variance
|
||||
|
||||
# Create 2D tile: [ROWS_PER_BLOCK, BLOCK_N]
|
||||
rows = row_start + tl.arange(0, ROWS_PER_BLOCK)
|
||||
cols = tl.arange(0, BLOCK_N)
|
||||
x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
|
||||
|
||||
# Compute offsets for 2D tile
|
||||
row_offsets = rows[:, None] * stride_x_row
|
||||
col_offsets = cols[None, :] + group * N
|
||||
|
||||
# Base pointers
|
||||
X_base = X + row_offsets + col_offsets
|
||||
Y_base = Y + rows[:, None] * stride_y_row + col_offsets
|
||||
|
||||
# Create mask for valid rows and columns
|
||||
row_mask = rows[:, None] < M
|
||||
col_mask = cols[None, :] < N
|
||||
mask = row_mask & col_mask
|
||||
|
||||
# Load input data with 2D tile
|
||||
x = tl.load(X_base, mask=mask, other=0.0).to(tl.float32)
|
||||
|
||||
if HAS_Z and not NORM_BEFORE_GATE:
|
||||
z = tl.load(Z + cols, mask=cols < N).to(tl.float32)
|
||||
Z_base = Z + rows[:, None] * stride_z_row + col_offsets
|
||||
z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
|
||||
x *= z * tl.sigmoid(z)
|
||||
|
||||
# Compute mean and variance per row (reduce along axis 1)
|
||||
if not IS_RMS_NORM:
|
||||
mean = tl.sum(x, axis=0) / N
|
||||
tl.store(Mean + row, mean)
|
||||
xbar = tl.where(cols < N, x - mean, 0.0)
|
||||
var = tl.sum(xbar * xbar, axis=0) / N
|
||||
mean = tl.sum(x, axis=1) / N # Shape: [ROWS_PER_BLOCK]
|
||||
# Store mean for each row
|
||||
mean_offsets = group * M + rows
|
||||
mean_mask = rows < M
|
||||
tl.store(Mean + mean_offsets, mean, mask=mean_mask)
|
||||
# Broadcast mean back to 2D for subtraction
|
||||
xbar = tl.where(mask, x - mean[:, None], 0.0)
|
||||
var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK]
|
||||
else:
|
||||
xbar = tl.where(cols < N, x, 0.0)
|
||||
var = tl.sum(xbar * xbar, axis=0) / N
|
||||
rstd = 1 / tl.sqrt(var + eps)
|
||||
tl.store(Rstd + row, rstd)
|
||||
# Normalize and apply linear transformation
|
||||
mask = cols < N
|
||||
w = tl.load(W + cols, mask=mask).to(tl.float32)
|
||||
xbar = tl.where(mask, x, 0.0)
|
||||
var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK]
|
||||
mean = 0.0 # Placeholder for RMS norm
|
||||
|
||||
rstd = tl.rsqrt(var + eps) # Shape: [ROWS_PER_BLOCK]
|
||||
|
||||
# Store rstd for each row
|
||||
rstd_offsets = group * M + rows
|
||||
rstd_mask = rows < M
|
||||
tl.store(Rstd + rstd_offsets, rstd, mask=rstd_mask)
|
||||
|
||||
# Load weights and biases (broadcast across rows)
|
||||
w_offsets = cols + group * N
|
||||
w_mask = cols < N
|
||||
w = tl.load(W + w_offsets, mask=w_mask, other=0.0).to(tl.float32)
|
||||
|
||||
if HAS_BIAS:
|
||||
b = tl.load(B + cols, mask=mask).to(tl.float32)
|
||||
x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
|
||||
y = x_hat * w + b if HAS_BIAS else x_hat * w
|
||||
b = tl.load(B + w_offsets, mask=w_mask, other=0.0).to(tl.float32)
|
||||
|
||||
# Normalize and apply linear transformation
|
||||
if not IS_RMS_NORM:
|
||||
x_hat = (x - mean[:, None]) * rstd[:, None]
|
||||
else:
|
||||
x_hat = x * rstd[:, None]
|
||||
|
||||
y = x_hat * w[None, :] + b[None, :] if HAS_BIAS else x_hat * w[None, :]
|
||||
|
||||
if HAS_Z and NORM_BEFORE_GATE:
|
||||
z = tl.load(Z + cols, mask=mask).to(tl.float32)
|
||||
Z_base = Z + rows[:, None] * stride_z_row + col_offsets
|
||||
z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
|
||||
y *= z * tl.sigmoid(z)
|
||||
|
||||
# Write output
|
||||
tl.store(Y + cols, y, mask=mask)
|
||||
tl.store(Y_base, y, mask=mask)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def _get_sm_count(device: torch.device) -> int:
|
||||
"""Get and cache the SM count for a given device."""
|
||||
props = torch.cuda.get_device_properties(device)
|
||||
return props.multi_processor_count
|
||||
|
||||
|
||||
def calc_rows_per_block(M: int, device: torch.device) -> int:
|
||||
sm_count = _get_sm_count(device)
|
||||
rows_per_block = next_power_of_2(cdiv(M, 2 * sm_count))
|
||||
rows_per_block = min(rows_per_block, 4)
|
||||
return rows_per_block
|
||||
|
||||
|
||||
def _layer_norm_fwd(
|
||||
@@ -158,7 +208,10 @@ def _layer_norm_fwd(
|
||||
raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
|
||||
# heuristics for number of warps
|
||||
num_warps = min(max(BLOCK_N // 256, 1), 8)
|
||||
grid = (M, ngroups)
|
||||
# Calculate rows per block based on SM count
|
||||
rows_per_block = calc_rows_per_block(M, x.device)
|
||||
# Update grid to use rows_per_block
|
||||
grid = (cdiv(M, rows_per_block), ngroups)
|
||||
with device_context(x.device):
|
||||
_layer_norm_fwd_1pass_kernel[grid](
|
||||
x,
|
||||
@@ -175,6 +228,7 @@ def _layer_norm_fwd(
|
||||
group_size,
|
||||
eps,
|
||||
BLOCK_N=BLOCK_N,
|
||||
ROWS_PER_BLOCK=rows_per_block,
|
||||
HAS_BIAS=bias is not None,
|
||||
HAS_Z=z is not None,
|
||||
NORM_BEFORE_GATE=norm_before_gate,
|
||||
|
||||
@@ -2467,6 +2467,11 @@ def set_cuda_arch():
|
||||
)
|
||||
|
||||
|
||||
def cdiv(a: int, b: int) -> int:
|
||||
"""Ceiling division."""
|
||||
return -(a // -b)
|
||||
|
||||
|
||||
def next_power_of_2(n: int):
|
||||
return 1 << (n - 1).bit_length() if n > 0 else 1
|
||||
|
||||
|
||||
@@ -70,6 +70,7 @@ suites = {
|
||||
TestFile("models/test_ministral3_models.py"),
|
||||
TestFile("test_mistral_large3_basic.py"),
|
||||
TestFile("test_prefill_delayer.py"),
|
||||
TestFile("test_fla_layernorm_guard.py"),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
384
test/srt/test_fla_layernorm_guard.py
Normal file
384
test/srt/test_fla_layernorm_guard.py
Normal file
@@ -0,0 +1,384 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import socket
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
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 layernorm_fn, rms_norm_ref
|
||||
|
||||
# Optional dependency in sglang repo; skip collection cleanly if absent.
|
||||
custom_all_reduce_utils = pytest.importorskip(
|
||||
"sglang.srt.distributed.device_communicators.custom_all_reduce_utils"
|
||||
)
|
||||
parallel_state = pytest.importorskip("sglang.srt.distributed.parallel_state")
|
||||
|
||||
update_environment_variables = custom_all_reduce_utils.update_environment_variables
|
||||
init_distributed_environment = parallel_state.init_distributed_environment
|
||||
initialize_model_parallel = parallel_state.initialize_model_parallel
|
||||
|
||||
NUM_GPUS = 2
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
# Avoid hard-coded port collisions when pytest runs tests in parallel.
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("localhost", 0))
|
||||
s.listen(1)
|
||||
return int(s.getsockname()[1])
|
||||
|
||||
|
||||
def _skip_if_no_cuda_or_not_enough_gpus(required_gpus: int = NUM_GPUS) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA device not available")
|
||||
if torch.cuda.device_count() < required_gpus:
|
||||
pytest.skip(f"Need >= {required_gpus} GPUs, got {torch.cuda.device_count()}")
|
||||
|
||||
|
||||
def _skip_if_dtype_unsupported(dtype: torch.dtype) -> None:
|
||||
if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported():
|
||||
pytest.skip("bfloat16 not supported on this CUDA device")
|
||||
|
||||
|
||||
def _setup_sglang_distributed(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
master_port: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.device:
|
||||
# Match sglang test style: set per-rank CUDA device + default dtype/device.
|
||||
torch.manual_seed(0)
|
||||
torch.cuda.manual_seed_all(0)
|
||||
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
if hasattr(torch, "set_default_device"):
|
||||
torch.set_default_device(device)
|
||||
if hasattr(torch, "set_default_dtype"):
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": str(master_port),
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment(
|
||||
world_size=world_size, rank=local_rank, local_rank=local_rank
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
return device
|
||||
|
||||
|
||||
def layer_norm_ref(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
z: torch.Tensor | None = None,
|
||||
eps: float = 1e-6,
|
||||
group_size: int | None = None,
|
||||
norm_before_gate: bool = True,
|
||||
is_rms_norm: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Reference implementation for both LayerNorm and RMSNorm (supports optional gate + group norm)."""
|
||||
if is_rms_norm:
|
||||
return rms_norm_ref(
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
z=z,
|
||||
eps=eps,
|
||||
group_size=group_size,
|
||||
norm_before_gate=norm_before_gate,
|
||||
upcast=True,
|
||||
)
|
||||
|
||||
dtype = x.dtype
|
||||
x_f = x.float()
|
||||
w_f = weight.float()
|
||||
b_f = bias.float() if bias is not None else None
|
||||
z_f = z.float() if z is not None else None
|
||||
|
||||
if z_f is not None and not norm_before_gate:
|
||||
x_f = x_f * F.silu(z_f)
|
||||
|
||||
if group_size is None:
|
||||
mean = x_f.mean(dim=-1, keepdim=True)
|
||||
var = (x_f - mean).square().mean(dim=-1, keepdim=True)
|
||||
rstd = torch.rsqrt(var + eps)
|
||||
out = (x_f - mean) * rstd * w_f
|
||||
if b_f is not None:
|
||||
out = out + b_f
|
||||
else:
|
||||
hidden = x_f.shape[-1]
|
||||
assert hidden % group_size == 0
|
||||
ng = hidden // group_size
|
||||
xg = x_f.view(*x_f.shape[:-1], ng, group_size)
|
||||
mean = xg.mean(dim=-1, keepdim=True)
|
||||
var = (xg - mean).square().mean(dim=-1, keepdim=True)
|
||||
rstd = torch.rsqrt(var + eps)
|
||||
xg = (xg - mean) * rstd
|
||||
out = xg.reshape(*x_f.shape[:-1], hidden) * w_f
|
||||
if b_f is not None:
|
||||
out = out + b_f
|
||||
|
||||
if z_f is not None and norm_before_gate:
|
||||
out = out * F.silu(z_f)
|
||||
|
||||
return out.to(dtype)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FwdCase:
|
||||
name: str
|
||||
with_gate: bool
|
||||
norm_before_gate: bool
|
||||
group_size: int | None
|
||||
is_rms_norm: bool
|
||||
|
||||
|
||||
CASES: list[FwdCase] = [
|
||||
FwdCase(
|
||||
"layernorm",
|
||||
with_gate=False,
|
||||
norm_before_gate=True,
|
||||
group_size=None,
|
||||
is_rms_norm=False,
|
||||
),
|
||||
FwdCase(
|
||||
"rmsnorm",
|
||||
with_gate=False,
|
||||
norm_before_gate=True,
|
||||
group_size=None,
|
||||
is_rms_norm=True,
|
||||
),
|
||||
FwdCase(
|
||||
"layernorm_gate_pre",
|
||||
with_gate=True,
|
||||
norm_before_gate=True,
|
||||
group_size=None,
|
||||
is_rms_norm=False,
|
||||
),
|
||||
FwdCase(
|
||||
"layernorm_gate_post",
|
||||
with_gate=True,
|
||||
norm_before_gate=False,
|
||||
group_size=None,
|
||||
is_rms_norm=False,
|
||||
),
|
||||
FwdCase(
|
||||
"rmsnorm_gate_pre",
|
||||
with_gate=True,
|
||||
norm_before_gate=True,
|
||||
group_size=None,
|
||||
is_rms_norm=True,
|
||||
),
|
||||
FwdCase(
|
||||
"group_layernorm",
|
||||
with_gate=False,
|
||||
norm_before_gate=True,
|
||||
group_size=128,
|
||||
is_rms_norm=False,
|
||||
),
|
||||
FwdCase(
|
||||
"group_rmsnorm",
|
||||
with_gate=False,
|
||||
norm_before_gate=True,
|
||||
group_size=128,
|
||||
is_rms_norm=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [128])
|
||||
@pytest.mark.parametrize("hidden_size", [256])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("case", CASES, ids=lambda c: c.name)
|
||||
def test_layernorm_guard_fwd_spawn(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
case: FwdCase,
|
||||
device: str = "cuda",
|
||||
):
|
||||
_skip_if_no_cuda_or_not_enough_gpus(NUM_GPUS)
|
||||
_skip_if_dtype_unsupported(dtype)
|
||||
|
||||
if case.group_size is not None and hidden_size % case.group_size != 0:
|
||||
pytest.skip(
|
||||
f"hidden_size {hidden_size} not divisible by group_size {case.group_size}"
|
||||
)
|
||||
|
||||
master_port = _find_free_port()
|
||||
world_size = NUM_GPUS
|
||||
|
||||
torch.multiprocessing.spawn(
|
||||
_layernorm_guard_fwd_worker,
|
||||
args=(
|
||||
world_size,
|
||||
master_port,
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype,
|
||||
case,
|
||||
device,
|
||||
),
|
||||
nprocs=world_size,
|
||||
join=True,
|
||||
)
|
||||
|
||||
|
||||
def _layernorm_guard_fwd_worker(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
master_port: int,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
case: FwdCase,
|
||||
device: str,
|
||||
):
|
||||
device = _setup_sglang_distributed(local_rank, world_size, master_port, dtype)
|
||||
|
||||
with torch.inference_mode():
|
||||
torch.manual_seed(42 + local_rank)
|
||||
torch.cuda.manual_seed_all(42 + local_rank)
|
||||
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
z = (
|
||||
torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
if case.with_gate
|
||||
else None
|
||||
)
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
bias = (
|
||||
None
|
||||
if case.is_rms_norm
|
||||
else torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
)
|
||||
eps = 1e-6
|
||||
|
||||
out, mean, rstd = layer_norm_fwd(
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
eps,
|
||||
z=z,
|
||||
group_size=case.group_size,
|
||||
norm_before_gate=case.norm_before_gate,
|
||||
is_rms_norm=case.is_rms_norm,
|
||||
)
|
||||
|
||||
ref_out = layer_norm_ref(
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
z=z,
|
||||
eps=eps,
|
||||
group_size=case.group_size,
|
||||
norm_before_gate=case.norm_before_gate,
|
||||
is_rms_norm=case.is_rms_norm,
|
||||
)
|
||||
|
||||
assert out.shape == x.shape
|
||||
assert out.dtype == x.dtype
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
|
||||
|
||||
# mean/rstd shape checks (same spirit as original vLLM tests)
|
||||
if case.group_size is None:
|
||||
if not case.is_rms_norm:
|
||||
assert mean.shape == (num_tokens,)
|
||||
assert rstd.shape == (num_tokens,)
|
||||
else:
|
||||
ngroups = hidden_size // case.group_size
|
||||
if not case.is_rms_norm:
|
||||
assert mean.shape == (ngroups * num_tokens,)
|
||||
assert rstd.shape == (ngroups * num_tokens,)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
def test_layernorm_guard_misc_spawn(dtype: torch.dtype, device: str = "cuda"):
|
||||
_skip_if_no_cuda_or_not_enough_gpus(NUM_GPUS)
|
||||
_skip_if_dtype_unsupported(dtype)
|
||||
|
||||
master_port = _find_free_port()
|
||||
world_size = NUM_GPUS
|
||||
|
||||
torch.multiprocessing.spawn(
|
||||
_layernorm_guard_misc_worker,
|
||||
args=(world_size, master_port, dtype, device),
|
||||
nprocs=world_size,
|
||||
join=True,
|
||||
)
|
||||
|
||||
|
||||
def _layernorm_guard_misc_worker(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
master_port: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
):
|
||||
device = _setup_sglang_distributed(local_rank, world_size, master_port, dtype)
|
||||
|
||||
with torch.inference_mode():
|
||||
torch.manual_seed(123 + local_rank)
|
||||
torch.cuda.manual_seed_all(123 + local_rank)
|
||||
|
||||
# 1) rows_per_block-like sizes
|
||||
hidden_size = 1024
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
bias = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
eps = 1e-6
|
||||
for num_tokens in [513]:
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
out, _, _ = layer_norm_fwd(x, weight, bias, eps, z=None, is_rms_norm=False)
|
||||
ref = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
|
||||
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
|
||||
|
||||
# 2) strided input (slice then contiguous)
|
||||
num_tokens = 128
|
||||
x_large = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
|
||||
x = x_large[:, :hidden_size]
|
||||
x_contig = x.contiguous()
|
||||
out, _, _ = layer_norm_fwd(
|
||||
x_contig, weight, bias, eps, z=None, is_rms_norm=False
|
||||
)
|
||||
ref = layer_norm_ref(x_contig, weight, bias, z=None, eps=eps, is_rms_norm=False)
|
||||
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
|
||||
|
||||
# 3) provided output buffer
|
||||
num_tokens = 256
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
out_buf = torch.empty_like(x)
|
||||
out, _, _ = layer_norm_fwd(
|
||||
x, weight, bias, eps, z=None, out=out_buf, is_rms_norm=False
|
||||
)
|
||||
assert out.data_ptr() == out_buf.data_ptr()
|
||||
ref = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
|
||||
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
|
||||
|
||||
# 4) multidimensional input via autograd fn
|
||||
for shape in [(4, 16, 1024)]:
|
||||
hs = shape[-1]
|
||||
x = torch.randn(*shape, dtype=dtype, device=device)
|
||||
w = torch.randn(hs, dtype=dtype, device=device)
|
||||
b = torch.randn(hs, dtype=dtype, device=device)
|
||||
out = layernorm_fn(x, w, b, z=None, eps=eps)
|
||||
ref = layer_norm_ref(x, w, b, z=None, eps=eps, is_rms_norm=False)
|
||||
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
Reference in New Issue
Block a user