diff --git a/benchmark/fla/benchmark_layernorm_gated.py b/benchmark/fla/benchmark_layernorm_gated.py new file mode 100644 index 000000000..82440582b --- /dev/null +++ b/benchmark/fla/benchmark_layernorm_gated.py @@ -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() diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index ed64cc438..5d55247da 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py +++ b/python/sglang/srt/layers/attention/fla/layernorm_gated.py @@ -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, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 8560246c6..4b2110ed8 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -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 diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index b0a6a38f1..ca0782127 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -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"), ], } diff --git a/test/srt/test_fla_layernorm_guard.py b/test/srt/test_fla_layernorm_guard.py new file mode 100644 index 000000000..96b26237f --- /dev/null +++ b/test/srt/test_fla_layernorm_guard.py @@ -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__])