diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index fa0230fb6..d4e71f535 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -373,6 +373,23 @@ class CustomAllreduce: ) return out + def deterministic_all_reduce( + self, + inp: torch.Tensor, + *, + out: torch.Tensor = None, + registered: bool = False, + ): + """Deterministic all-reduce using 1-stage kernel with fixed ordering (AMD only).""" + if out is None: + out = torch.empty_like(inp) + if registered: + ops.deterministic_all_reduce_reg(self._ptr, inp, out) + else: + reg_buffer = self.buffer.view(inp.dtype)[: inp.numel()] + ops.deterministic_all_reduce_unreg(self._ptr, inp, reg_buffer, out) + return out + def custom_all_reduce(self, input: torch.Tensor) -> Optional[torch.Tensor]: """The main allreduce API that provides support for cuda graph.""" # When custom allreduce is disabled, this will be None. @@ -411,20 +428,37 @@ class CustomAllreduce: def dispatch_custom_allreduce(): - """Return the CustomAllreduce class to use (aiter on ROCm if enabled).""" + """Return the CustomAllreduce class to use (aiter on ROCm if enabled). + + On AMD with 1-stage AR enabled, use sglang's CustomAllreduce (has deterministic_all_reduce method). + Otherwise use AiterCustomAllreduce if available. + """ + # Check if 1-stage AR should be used + if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set(): + use_1stage = envs.SGLANG_USE_1STAGE_ALLREDUCE.get() + else: + use_1stage = envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() + + # On AMD with 1-stage AR, use sglang's CustomAllreduce + # (AiterCustomAllreduce doesn't have deterministic_all_reduce method) + if is_hip() and use_1stage: + logger.info("[AR] Using sglang CustomAllreduce (1-stage kernel)") + return CustomAllreduce + if is_hip() and get_bool_env_var("SGLANG_USE_AITER_AR", default="true"): try: from aiter.dist.device_communicators.custom_all_reduce import ( CustomAllreduce as AiterCustomAllreduce, ) - logger.info("Using AiterCustomAllreduce for ROCm.") + logger.info("[AR] Using AiterCustomAllreduce (AMD default)") return AiterCustomAllreduce except ImportError as e: logger.warning( - "Aiter custom all-reduce not available (optional dependency missing); " + "[AR] Aiter custom all-reduce not available; " "falling back to sglang CustomAllreduce. Details: %s", e, ) return CustomAllreduce + logger.info("[AR] Using sglang CustomAllreduce") return CustomAllreduce diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py index 3d7e3a56b..c312e8f2c 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py @@ -90,6 +90,16 @@ elif _is_hip: ) -> None: _custom_ar.all_reduce_unreg(fa, inp, reg_buffer, out) + def deterministic_all_reduce_reg( + fa: int, inp: torch.Tensor, out: torch.Tensor + ) -> None: + _custom_ar.deterministic_all_reduce_reg(fa, inp, out) + + def deterministic_all_reduce_unreg( + fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor + ) -> None: + _custom_ar.deterministic_all_reduce_unreg(fa, inp, reg_buffer, out) + def dispose(fa: int) -> None: _custom_ar.dispose(fa) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 216b8d5a3..7709cc926 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -376,6 +376,23 @@ class GroupCoordinator: group=self.cpu_group, device=self.device, ) + # Log which all-reduce mode will be used + if is_hip(): + if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set(): + if envs.SGLANG_USE_1STAGE_ALLREDUCE.get(): + logger.info( + "[AR] All-reduce: 1-stage kernel (SGLANG_USE_1STAGE_ALLREDUCE=1)" + ) + else: + logger.info( + "[AR] All-reduce: default (SGLANG_USE_1STAGE_ALLREDUCE=0)" + ) + elif envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get(): + logger.info( + "[AR] All-reduce: 1-stage kernel (deterministic inference enabled)" + ) + else: + logger.info("[AR] All-reduce: default") except Exception as e: logger.warning( f"Setup Custom allreduce failed with {e}. To silence this " @@ -393,6 +410,8 @@ class GroupCoordinator: ) except Exception as e: logger.warning(f"Failed to initialize QuickAllReduce: {e}") + elif self.world_size > 1 and is_hip(): + logger.info("[AR] All-reduce call path: NCCL (custom AR disabled)") self.torch_symm_mem_comm: Optional[TorchSymmMemCommunicator] = None if self.use_torch_symm_mem_all_reduce and self.world_size > 1: @@ -560,6 +579,26 @@ class GroupCoordinator: if self.world_size == 1: return input_ + # On AMD, use the deterministic 1-stage kernel when: + # - SGLANG_USE_1STAGE_ALLREDUCE=1 (explicitly enabled), OR + # - SGLANG_USE_1STAGE_ALLREDUCE not set AND --enable-deterministic-inference is on + if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set(): + use_1stage_ar = envs.SGLANG_USE_1STAGE_ALLREDUCE.get() + else: + use_1stage_ar = envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() + use_deterministic_ar = is_hip() and use_1stage_ar + if use_deterministic_ar: + if not input_.is_cpu and self.ca_comm is not None: + inp_size = input_.numel() * input_.element_size() + # Try unregistered mode first (faster for smaller tensors) + if inp_size < self.ca_comm.max_size: + return self.ca_comm.deterministic_all_reduce( + input_, registered=False + ) + # Use registered mode for larger tensors + self.ca_comm.register_buffer(input_) + return self.ca_comm.deterministic_all_reduce(input_, registered=True) + if input_.is_cpu: if is_shm_available(input_.dtype, self.world_size, self.local_size): torch.ops.sgl_kernel.shm_allreduce(input_, REDUCE_OP_SUM) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index bfc606375..9fceae6fe 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -308,6 +308,11 @@ class Envs: # Deterministic inference SGLANG_ENABLE_DETERMINISTIC_INFERENCE = EnvBool(False) + # Use 1-stage all-reduce kernel on AMD (deterministic, fixed accumulation order) + # If not set: auto (enabled when --enable-deterministic-inference is on) + # Set to 1: force enable (even without --enable-deterministic-inference) + # Set to 0: force disable (use default Aiter AR even with --enable-deterministic-inference) + SGLANG_USE_1STAGE_ALLREDUCE = EnvBool(False) SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096) SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048) SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e54f1c9a2..56ac53a5d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2314,11 +2314,19 @@ class ServerArgs: # Check TP size if self.tp_size > 1: - os.environ["NCCL_ALGO"] = "allreduce:tree" - self.disable_custom_all_reduce = True - logger.warning( - "NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when TP size > 1." - ) + if is_hip(): + # AMD: use 1-stage all-reduce kernel which is inherently deterministic + # (each GPU reads all data from all GPUs, reduces locally in fixed order) + logger.info( + "AMD/ROCm: Using 1-stage all-reduce kernel (deterministic)" + ) + else: + # CUDA: use NCCL tree algorithm + os.environ["NCCL_ALGO"] = "allreduce:tree" + self.disable_custom_all_reduce = True + logger.warning( + "NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when TP size > 1." + ) def _handle_dllm_inference(self): if self.dllm_algorithm is None: diff --git a/sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py b/sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py new file mode 100644 index 000000000..647543089 --- /dev/null +++ b/sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py @@ -0,0 +1,692 @@ +""" +Benchmark latency comparison between different all-reduce implementations. + +Compares: +- NCCL all-reduce (may be non-deterministic) +- Reduce-scatter + all-gather (RS+AG, deterministic but slower) +- Deterministic 1-stage kernel (forces fixed accumulation order, deterministic) + +Note: The "deterministic kernel" is NOT RS+AG. It uses the 1-stage kernel where +each GPU reads all data from all GPUs and reduces locally in a fixed order. + +Usage: + python bench_amd_deterministic_allreduce.py +""" + +import multiprocessing as mp +import os +import socket +import statistics +import sys +import time + +import torch +import torch.distributed as dist + +# Add python directory to path to import sglang modules +script_dir = os.path.dirname(os.path.abspath(__file__)) +python_dir = os.path.join(script_dir, "python") +sys.path.insert(0, python_dir) + +# Try to import custom all-reduce if available +try: + import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as custom_ar_ops + from sglang.srt.distributed.device_communicators.custom_all_reduce import ( + CustomAllreduce, + ) + from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import ( + is_weak_contiguous, + ) + + CUSTOM_AR_AVAILABLE = custom_ar_ops.IS_CUSTOM_AR_AVAILABLE +except (ImportError, AttributeError): + CUSTOM_AR_AVAILABLE = False + CustomAllreduce = None + is_weak_contiguous = None + +# Note: sglang's optimized all-reduce requires full runtime initialization +# and won't work in standalone benchmarks, so we skip it +SGLANG_AVAILABLE = False + + +def get_open_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def init_custom_ar_if_available(rank, world_size, device): + """Check if custom all-reduce is available and applicable.""" + if not CUSTOM_AR_AVAILABLE or CustomAllreduce is None: + return False + + # Custom AR works best for single-node, even number of GPUs, world_size <= 8 + if world_size <= 8 and world_size % 2 == 0: + return True + + return False + + +def reduce_scatter_then_all_gather(tensor, rank, world_size, custom_ar=None): + """ + Deterministic all-reduce using reduce-scatter + all-gather. + This is deterministic because it uses fixed ordering (no atomics). + """ + total_size = tensor.numel() + if total_size % world_size != 0: + # Fallback to all-gather + local reduce if not divisible + gather_list = [torch.empty_like(tensor) for _ in range(world_size)] + dist.all_gather(gather_list, tensor) + stacked = torch.stack(gather_list, dim=0) + tensor.copy_(stacked.sum(dim=0)) + return + + chunk_size = total_size // world_size + + # Flatten to 1D + tensor_flat = tensor.view(-1) + + # Reduce-scatter: each rank gets its chunk of the reduced result + output_chunk = torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device) + + # Split input into chunks for reduce-scatter + input_chunks = [ + tensor_flat[i * chunk_size : (i + 1) * chunk_size].clone() + for i in range(world_size) + ] + + dist.reduce_scatter(output_chunk, input_chunks) + + # All-gather: broadcast each rank's chunk to all ranks + output_chunks = [ + torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device) + for _ in range(world_size) + ] + dist.all_gather(output_chunks, output_chunk) + + # Concatenate results back + result_flat = torch.cat(output_chunks, dim=0) + tensor.copy_(result_flat.view(tensor.shape)) + + +def worker(world_size, rank, port, results_queue): + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + ) + + # Try to initialize custom all-reduce if available + custom_ar = None + use_custom_ar = init_custom_ar_if_available(rank, world_size, device) + if use_custom_ar and CUSTOM_AR_AVAILABLE: + try: + # Create a gloo group for custom AR (it requires non-NCCL backend) + # All ranks must call new_group with the same parameters + from torch.distributed import new_group + + dist.barrier() # Ensure all ranks are ready + ar_group = new_group(backend="gloo") + dist.barrier() # Ensure group creation is complete + custom_ar = CustomAllreduce(group=ar_group, device=device) + if rank == 0: + print(" Using custom all-reduce (deterministic)") + except Exception as e: + if rank == 0: + print(f" Custom AR init failed: {e}, using NCCL fallback") + custom_ar = None + dist.barrier() # Ensure all ranks continue even if one fails + + # Test different batch sizes - similar to test_ar.py + batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256, 512] + hidden_dim = 16384 # Fixed hidden dimension + + num_trials = 10 # Same as test_ar.py + + # Different seed per rank - each GPU has DIFFERENT input (like test_ar.py) + torch.manual_seed(42 + rank) + + results = {} + + for bs in batch_sizes: + # Create fixed input for all trials (like test_ar.py) + base_input = torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) + + dist.barrier() + + if rank == 0: + print(f"\nBatch size {bs:4d}:") + print(f" Testing determinism across {num_trials} trials...") + + # Test all-reduce determinism + results_ar = [] + latencies_ar = [] + for trial in range(num_trials): + # Clone the same input for each trial + inp_ar = base_input.clone() + inp_flat_ar = inp_ar.view(-1) + + # Measure latency + torch.cuda.synchronize() + start = time.perf_counter() + dist.all_reduce(inp_flat_ar, op=dist.ReduceOp.SUM) + torch.cuda.synchronize() + end = time.perf_counter() + latencies_ar.append(end - start) + + # Store checksum and first values (like test_ar.py) + checksum = inp_flat_ar.sum().item() + first_vals = inp_flat_ar[:5].clone() + results_ar.append((checksum, first_vals)) + + # Test reduce-scatter + all-gather determinism + results_rs_ag = [] + latencies_rs_ag = [] + for trial in range(num_trials): + # Clone the same input for each trial + inp_rs_ag = base_input.clone() + inp_flat_rs_ag = inp_rs_ag.view(-1) + + # Measure latency + torch.cuda.synchronize() + start = time.perf_counter() + reduce_scatter_then_all_gather( + inp_flat_rs_ag, rank, world_size, custom_ar=None + ) + torch.cuda.synchronize() + end = time.perf_counter() + latencies_rs_ag.append(end - start) + + # Store checksum and first values (like test_ar.py) + checksum = inp_flat_rs_ag.sum().item() + first_vals = inp_flat_rs_ag[:5].clone() + results_rs_ag.append((checksum, first_vals)) + + # Note: sglang's optimized all-reduce requires full runtime initialization + # and is not tested in this standalone benchmark + use_sglang_optimized = False + results_optimized_rs_ag = [] + latencies_optimized_rs_ag = [] + + # Test custom all-reduce determinism (if available) + results_custom_ar = [] + latencies_custom_ar = [] + if custom_ar is not None: + for trial in range(num_trials): + # Clone the same input for each trial + inp_custom = base_input.clone() + inp_flat_custom = inp_custom.view(-1) + + # Measure latency + torch.cuda.synchronize() + start = time.perf_counter() + reduce_scatter_then_all_gather( + inp_flat_custom, rank, world_size, custom_ar=custom_ar + ) + torch.cuda.synchronize() + end = time.perf_counter() + latencies_custom_ar.append(end - start) + + # Store checksum and first values (like test_ar.py) + checksum = inp_flat_custom.sum().item() + first_vals = inp_flat_custom[:5].clone() + results_custom_ar.append((checksum, first_vals)) + + # Test deterministic kernel (if available) + results_deterministic_kernel = [] + latencies_deterministic_kernel = [] + deterministic_kernel_available = False + if custom_ar is not None and hasattr(custom_ar, "deterministic_all_reduce"): + # Check if input size fits in buffer + input_size_bytes = base_input.numel() * base_input.element_size() + if input_size_bytes > custom_ar.max_size: + if rank == 0: + print( + f" Deterministic kernel skipped: input size ({input_size_bytes/(1024*1024):.1f} MB) > buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)" + ) + deterministic_kernel_available = False + else: + try: + deterministic_kernel_available = True + for trial in range(num_trials): + # Clone the same input for each trial + inp_kernel = base_input.clone() + + # Measure latency + torch.cuda.synchronize() + start = time.perf_counter() + result_kernel = custom_ar.deterministic_all_reduce( + inp_kernel, registered=False + ) + torch.cuda.synchronize() + end = time.perf_counter() + latencies_deterministic_kernel.append(end - start) + + # Store checksum and first values + result_flat_kernel = result_kernel.view(-1) + checksum = result_flat_kernel.sum().item() + first_vals = result_flat_kernel[:5].clone() + results_deterministic_kernel.append((checksum, first_vals)) + except Exception as e: + if rank == 0: + print( + f" Deterministic kernel test failed for batch size {bs}: {e}" + ) + deterministic_kernel_available = False + + dist.barrier() + + if rank == 0: + # Check determinism for all-reduce + ar_deterministic = True + ar_ref_sum, ar_ref_vals = results_ar[0] + ar_variance = [] + for i, (s, vals) in enumerate(results_ar[1:], 1): + if abs(ar_ref_sum - s) > 1e-3 or not torch.allclose( + ar_ref_vals, vals, rtol=1e-3 + ): + ar_deterministic = False + ar_variance.append(abs(ar_ref_sum - s)) + + # Check determinism for reduce-scatter + all-gather + rs_ag_deterministic = True + rs_ag_ref_sum, rs_ag_ref_vals = results_rs_ag[0] + rs_ag_variance = [] + for i, (s, vals) in enumerate(results_rs_ag[1:], 1): + if abs(rs_ag_ref_sum - s) > 1e-3 or not torch.allclose( + rs_ag_ref_vals, vals, rtol=1e-3 + ): + rs_ag_deterministic = False + rs_ag_variance.append(abs(rs_ag_ref_sum - s)) + + # Check determinism for optimized RS+AG (if available) + optimized_rs_ag_deterministic = None + optimized_rs_ag_max_variance = None + lat_optimized_rs_ag_median = None + if use_sglang_optimized and results_optimized_rs_ag: + optimized_rs_ag_deterministic = True + opt_rs_ag_ref_sum, opt_rs_ag_ref_vals = results_optimized_rs_ag[0] + opt_rs_ag_variance = [] + for i, (s, vals) in enumerate(results_optimized_rs_ag[1:], 1): + if abs(opt_rs_ag_ref_sum - s) > 1e-3 or not torch.allclose( + opt_rs_ag_ref_vals, vals, rtol=1e-3 + ): + optimized_rs_ag_deterministic = False + opt_rs_ag_variance.append(abs(opt_rs_ag_ref_sum - s)) + optimized_rs_ag_max_variance = ( + max(opt_rs_ag_variance) if opt_rs_ag_variance else 0.0 + ) + lat_optimized_rs_ag_median = statistics.median( + latencies_optimized_rs_ag + ) + + # Check determinism for custom all-reduce (if available) + custom_ar_deterministic = None + custom_ar_max_variance = None + lat_custom_ar_median = None + if custom_ar is not None and results_custom_ar: + custom_ar_deterministic = True + custom_ar_ref_sum, custom_ar_ref_vals = results_custom_ar[0] + custom_ar_variance = [] + for i, (s, vals) in enumerate(results_custom_ar[1:], 1): + if abs(custom_ar_ref_sum - s) > 1e-3 or not torch.allclose( + custom_ar_ref_vals, vals, rtol=1e-3 + ): + custom_ar_deterministic = False + custom_ar_variance.append(abs(custom_ar_ref_sum - s)) + custom_ar_max_variance = ( + max(custom_ar_variance) if custom_ar_variance else 0.0 + ) + lat_custom_ar_median = statistics.median(latencies_custom_ar) + + # Check determinism for deterministic kernel (if available) + deterministic_kernel_deterministic = None + deterministic_kernel_max_variance = None + lat_deterministic_kernel_median = None + if deterministic_kernel_available and results_deterministic_kernel: + deterministic_kernel_deterministic = True + kernel_ref_sum, kernel_ref_vals = results_deterministic_kernel[0] + kernel_variance = [] + for i, (s, vals) in enumerate(results_deterministic_kernel[1:], 1): + if abs(kernel_ref_sum - s) > 1e-3 or not torch.allclose( + kernel_ref_vals, vals, rtol=1e-3 + ): + deterministic_kernel_deterministic = False + kernel_variance.append(abs(kernel_ref_sum - s)) + deterministic_kernel_max_variance = ( + max(kernel_variance) if kernel_variance else 0.0 + ) + lat_deterministic_kernel_median = statistics.median( + latencies_deterministic_kernel + ) + + # Calculate latency statistics + lat_ar_median = statistics.median(latencies_ar) + lat_rs_ag_median = statistics.median(latencies_rs_ag) + overhead_rs_ag = ((lat_rs_ag_median - lat_ar_median) / lat_ar_median) * 100 + + # Calculate variance statistics + ar_max_variance = max(ar_variance) if ar_variance else 0.0 + rs_ag_max_variance = max(rs_ag_variance) if rs_ag_variance else 0.0 + + results[bs] = { + "all_reduce": { + "latency_median": lat_ar_median, + "deterministic": ar_deterministic, + "max_variance": ar_max_variance, + }, + "rs_ag": { + "latency_median": lat_rs_ag_median, + "deterministic": rs_ag_deterministic, + "max_variance": rs_ag_max_variance, + }, + "custom_ar": ( + { + "latency_median": lat_custom_ar_median, + "deterministic": custom_ar_deterministic, + "max_variance": custom_ar_max_variance, + } + if custom_ar is not None + else None + ), + "deterministic_kernel": ( + { + "latency_median": lat_deterministic_kernel_median, + "deterministic": deterministic_kernel_deterministic, + "max_variance": deterministic_kernel_max_variance, + } + if lat_deterministic_kernel_median is not None + else None + ), + "optimized_rs_ag": ( + { + "latency_median": lat_optimized_rs_ag_median, + "deterministic": optimized_rs_ag_deterministic, + "max_variance": optimized_rs_ag_max_variance, + } + if lat_optimized_rs_ag_median is not None + else None + ), + "overhead_rs_ag_pct": overhead_rs_ag, + } + + print( + f" All-Reduce: {lat_ar_median*1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}" + ) + print( + f" RS+All-Gather: {lat_rs_ag_median*1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}" + ) + if custom_ar is not None and lat_custom_ar_median is not None: + overhead_custom = ( + (lat_custom_ar_median - lat_ar_median) / lat_ar_median + ) * 100 + print( + f" Custom AR: {lat_custom_ar_median*1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%" + ) + if lat_deterministic_kernel_median is not None: + overhead_kernel = ( + (lat_deterministic_kernel_median - lat_ar_median) / lat_ar_median + ) * 100 + speedup_kernel_vs_rs_ag = ( + (lat_rs_ag_median - lat_deterministic_kernel_median) + / lat_rs_ag_median + ) * 100 + print( + f" Deterministic Kernel: {lat_deterministic_kernel_median*1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%" + ) + if lat_optimized_rs_ag_median is not None: + overhead_opt = ( + (lat_optimized_rs_ag_median - lat_ar_median) / lat_ar_median + ) * 100 + speedup_vs_rs_ag = ( + (lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median + ) * 100 + print( + f" Optimized RS+AG: {lat_optimized_rs_ag_median*1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%" + ) + print(f" RS+AG Overhead: {overhead_rs_ag:+.1f}%") + + if rank == 0: + results_queue.put(results) + + dist.destroy_process_group() + + +def main(): + world_size = 8 + available_gpus = torch.cuda.device_count() + + print("=" * 80) + print("All-Reduce vs Reduce-Scatter + All-Gather Determinism & Latency Benchmark") + print("=" * 80) + print(f"Available GPUs: {available_gpus}") + print(f"Using world_size: {world_size}") + print(f"Hidden dimension: 16384") + print(f"Tensor dtype: bfloat16") + print(f"Trials per batch size: 10 (testing determinism)") + print(f"Testing batch sizes: [1, 4, 8, 16, 32, 64, 128, 256, 512]") + print("=" * 80) + + if available_gpus < world_size: + print( + f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead" + ) + world_size = available_gpus + + if world_size < 2: + print("ERROR: Need at least 2 GPUs for this benchmark") + return + + mp.set_start_method("spawn", force=True) + port = get_open_port() + + results_queue = mp.Queue() + procs = [] + for rank in range(world_size): + p = mp.Process(target=worker, args=(world_size, rank, port, results_queue)) + p.start() + procs.append(p) + + for p in procs: + p.join() + + # Collect results + if not results_queue.empty(): + results = results_queue.get() + + print("\n" + "=" * 80) + print("SUMMARY") + print("=" * 80) + header = f"{'Batch':<8} {'AR (ms)':<12} {'AR Det':<8} {'RS+AG (ms)':<15} {'RS+AG Det':<10} {'RS+AG Ovh':<12}" + if any(r.get("custom_ar") is not None for r in results.values()): + header += ( + f" {'Custom AR (ms)':<18} {'Custom AR Det':<15} {'Custom AR Ovh':<15}" + ) + if any(r.get("deterministic_kernel") is not None for r in results.values()): + header += f" {'Det Kernel (ms)':<18} {'Det Kernel Det':<15} {'Det Kernel Ovh':<15} {'Speedup':<10}" + if any(r.get("optimized_rs_ag") is not None for r in results.values()): + header += f" {'Opt RS+AG (ms)':<18} {'Opt RS+AG Det':<15} {'Opt RS+AG Ovh':<15} {'Speedup':<10}" + print(header) + print("-" * 150) + + for bs in sorted(results.keys()): + r = results[bs] + ar_det_str = "✓" if r["all_reduce"]["deterministic"] else "✗" + rs_ag_det_str = "✓" if r["rs_ag"]["deterministic"] else "✗" + line = ( + f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} " + f"{r['rs_ag']['latency_median']*1000:<15.3f} {rs_ag_det_str:<10} " + f"{r['overhead_rs_ag_pct']:<12.1f}" + ) + if r.get("custom_ar") is not None: + custom_ar = r["custom_ar"] + custom_ar_det_str = "✓" if custom_ar["deterministic"] else "✗" + custom_ar_overhead = ( + (custom_ar["latency_median"] - r["all_reduce"]["latency_median"]) + / r["all_reduce"]["latency_median"] + ) * 100 + line += f" {custom_ar['latency_median']*1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}" + if r.get("deterministic_kernel") is not None: + det_kernel = r["deterministic_kernel"] + det_kernel_det_str = "✓" if det_kernel["deterministic"] else "✗" + det_kernel_overhead = ( + (det_kernel["latency_median"] - r["all_reduce"]["latency_median"]) + / r["all_reduce"]["latency_median"] + ) * 100 + speedup_kernel = ( + (r["rs_ag"]["latency_median"] - det_kernel["latency_median"]) + / r["rs_ag"]["latency_median"] + ) * 100 + line += f" {det_kernel['latency_median']*1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}" + if r.get("optimized_rs_ag") is not None: + opt_rs_ag = r["optimized_rs_ag"] + opt_rs_ag_det_str = "✓" if opt_rs_ag["deterministic"] else "✗" + opt_rs_ag_overhead = ( + (opt_rs_ag["latency_median"] - r["all_reduce"]["latency_median"]) + / r["all_reduce"]["latency_median"] + ) * 100 + speedup = ( + (r["rs_ag"]["latency_median"] - opt_rs_ag["latency_median"]) + / r["rs_ag"]["latency_median"] + ) * 100 + line += f" {opt_rs_ag['latency_median']*1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}" + print(line) + + print("=" * 80) + + # Calculate statistics + overheads_rs_ag = [r["overhead_rs_ag_pct"] for r in results.values()] + ar_deterministic_count = sum( + 1 for r in results.values() if r["all_reduce"]["deterministic"] + ) + rs_ag_deterministic_count = sum( + 1 for r in results.values() if r["rs_ag"]["deterministic"] + ) + custom_ar_deterministic_count = sum( + 1 + for r in results.values() + if r.get("custom_ar") and r["custom_ar"]["deterministic"] + ) + custom_ar_total_count = sum( + 1 for r in results.values() if r.get("custom_ar") is not None + ) + + deterministic_kernel_deterministic_count = sum( + 1 + for r in results.values() + if r.get("deterministic_kernel") + and r["deterministic_kernel"]["deterministic"] + ) + deterministic_kernel_total_count = sum( + 1 for r in results.values() if r.get("deterministic_kernel") is not None + ) + + print(f"\nDeterminism Summary:") + print( + f" All-Reduce deterministic: {ar_deterministic_count}/{len(results)} batch sizes" + ) + print( + f" RS+All-Gather deterministic: {rs_ag_deterministic_count}/{len(results)} batch sizes" + ) + if custom_ar_total_count > 0: + print( + f" Custom AR deterministic: {custom_ar_deterministic_count}/{custom_ar_total_count} batch sizes" + ) + if deterministic_kernel_total_count > 0: + print( + f" Deterministic Kernel deterministic: {deterministic_kernel_deterministic_count}/{deterministic_kernel_total_count} batch sizes" + ) + + print(f"\nLatency Overhead Statistics (RS+AG vs All-Reduce):") + avg_overhead = statistics.mean(overheads_rs_ag) + median_overhead = statistics.median(overheads_rs_ag) + min_overhead = min(overheads_rs_ag) + max_overhead = max(overheads_rs_ag) + print(f" Average: {avg_overhead:.1f}%") + print(f" Median: {median_overhead:.1f}%") + print(f" Min: {min_overhead:.1f}%") + print(f" Max: {max_overhead:.1f}%") + + if custom_ar_total_count > 0: + overheads_custom = [] + for r in results.values(): + if r.get("custom_ar") is not None: + overhead = ( + ( + r["custom_ar"]["latency_median"] + - r["all_reduce"]["latency_median"] + ) + / r["all_reduce"]["latency_median"] + ) * 100 + overheads_custom.append(overhead) + print(f"\nLatency Overhead Statistics (Custom AR vs All-Reduce):") + print(f" Average: {statistics.mean(overheads_custom):.1f}%") + print(f" Median: {statistics.median(overheads_custom):.1f}%") + print(f" Min: {min(overheads_custom):.1f}%") + print(f" Max: {max(overheads_custom):.1f}%") + + if deterministic_kernel_total_count > 0: + overheads_kernel = [] + speedups_kernel = [] + for r in results.values(): + if r.get("deterministic_kernel") is not None: + overhead = ( + ( + r["deterministic_kernel"]["latency_median"] + - r["all_reduce"]["latency_median"] + ) + / r["all_reduce"]["latency_median"] + ) * 100 + overheads_kernel.append(overhead) + speedup = ( + ( + r["rs_ag"]["latency_median"] + - r["deterministic_kernel"]["latency_median"] + ) + / r["rs_ag"]["latency_median"] + ) * 100 + speedups_kernel.append(speedup) + print( + f"\nLatency Overhead Statistics (Deterministic Kernel vs All-Reduce):" + ) + print(f" Average: {statistics.mean(overheads_kernel):.1f}%") + print(f" Median: {statistics.median(overheads_kernel):.1f}%") + print(f" Min: {min(overheads_kernel):.1f}%") + print(f" Max: {max(overheads_kernel):.1f}%") + print(f"\nSpeedup Statistics (Deterministic Kernel vs RS+AG):") + print(f" Average: {statistics.mean(speedups_kernel):.1f}%") + print(f" Median: {statistics.median(speedups_kernel):.1f}%") + print(f" Min: {min(speedups_kernel):.1f}%") + print(f" Max: {max(speedups_kernel):.1f}%") + + # Show variance for non-deterministic cases + print(f"\nVariance Analysis (non-deterministic cases):") + for bs in sorted(results.keys()): + r = results[bs] + if not r["all_reduce"]["deterministic"]: + print( + f" Batch {bs}: All-Reduce max variance: {r['all_reduce']['max_variance']:.6f}" + ) + if not r["rs_ag"]["deterministic"]: + print( + f" Batch {bs}: RS+All-Gather max variance: {r['rs_ag']['max_variance']:.6f}" + ) + if r.get("custom_ar") is not None and not r["custom_ar"]["deterministic"]: + print( + f" Batch {bs}: Custom AR max variance: {r['custom_ar']['max_variance']:.6f}" + ) + if ( + r.get("deterministic_kernel") is not None + and not r["deterministic_kernel"]["deterministic"] + ): + print( + f" Batch {bs}: Deterministic Kernel max variance: {r['deterministic_kernel']['max_variance']:.6f}" + ) + + +if __name__ == "__main__": + main() diff --git a/sgl-kernel/csrc/allreduce/deterministic_all_reduce.hip b/sgl-kernel/csrc/allreduce/deterministic_all_reduce.hip new file mode 100644 index 000000000..8d78b97f4 --- /dev/null +++ b/sgl-kernel/csrc/allreduce/deterministic_all_reduce.hip @@ -0,0 +1,178 @@ +// Deterministic All-Reduce for ROCm/HIP +// +// This is a wrapper that forces the use of the existing 1-stage all-reduce kernel +// (cross_device_reduce_1stage) which is inherently deterministic due to fixed +// accumulation ordering (no atomics, no race conditions). +// +// How the 1-stage kernel works: +// - Each GPU reads ALL data from ALL other GPUs via direct memory access +// - Each GPU reduces the data locally in a fixed order +// - Result: every GPU has the complete reduced output +// +// This is NOT a reduce-scatter + all-gather (RS+AG) approach. +// The 2-stage kernel (cross_device_reduce_2stage) implements RS+AG but may have +// non-deterministic behavior, so we explicitly avoid it here. + +#include +#include +#include +#include + +#include "custom_all_reduce_hip.cuh" + +using fptr_t = int64_t; +static_assert(sizeof(void*) == sizeof(fptr_t)); + +// Helper function for weak contiguity check +bool _is_weak_contiguous_det(torch::Tensor& t) { + return t.is_contiguous() || + (t.storage().nbytes() - t.storage_offset() * t.element_size() == t.numel() * t.element_size()); +} + +// Deterministic all-reduce for registered buffers (ROCm) +// Uses the 1-stage kernel which is deterministic (fixed ordering) +void deterministic_all_reduce_reg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out) { + const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp)); + auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream(); + TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type()); + TORCH_CHECK_EQ(inp.numel(), out.numel()); + TORCH_CHECK(_is_weak_contiguous_det(out)); + TORCH_CHECK(_is_weak_contiguous_det(inp)); + + auto fa = reinterpret_cast(_fa); + + // For ROCm, manually call the 1-stage kernel to ensure deterministic ordering + // Get rank data pointer + sglang::RankData* ptrs; + hipStreamCaptureStatus status; + AT_CUDA_CHECK(hipStreamIsCapturing(stream, &status)); + if (status == hipStreamCaptureStatusActive) { + ptrs = fa->d_rank_data_base_ + fa->graph_unreg_buffers_.size(); + fa->graph_unreg_buffers_.push_back(inp.data_ptr()); + } else { + auto it = fa->buffers_.find(inp.data_ptr()); + if (it == fa->buffers_.end()) { + throw std::runtime_error("buffer not registered!"); + } + ptrs = it->second; + } + + int size = out.numel(); + int threads = 512; + + switch (out.scalar_type()) { + case at::ScalarType::Float: { + using T = float; + using P = typename sglang::packed_t::P; + auto d = P::size; + if (size % d != 0) { + throw std::runtime_error("size must be multiple of " + std::to_string(d)); + } + size /= d; + int blocks = std::min(16, (size + threads - 1) / threads); + // Always use 1-stage kernel for determinism + switch (fa->world_size_) { + case 2: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 4: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 6: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 8: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + default: + throw std::runtime_error("world_size must be in (2,4,6,8)"); + } + break; + } + case at::ScalarType::Half: { + using T = half; + using P = typename sglang::packed_t::P; + auto d = P::size; + if (size % d != 0) { + throw std::runtime_error("size must be multiple of " + std::to_string(d)); + } + size /= d; + int blocks = std::min(16, (size + threads - 1) / threads); + switch (fa->world_size_) { + case 2: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 4: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 6: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 8: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + default: + throw std::runtime_error("world_size must be in (2,4,6,8)"); + } + break; + } +#if (__HIP_ARCH__ >= 800 || !defined(__HIP_ARCH__)) + case at::ScalarType::BFloat16: { + using T = nv_bfloat16; + using P = typename sglang::packed_t::P; + auto d = P::size; + if (size % d != 0) { + throw std::runtime_error("size must be multiple of " + std::to_string(d)); + } + size /= d; + int blocks = std::min(16, (size + threads - 1) / threads); + switch (fa->world_size_) { + case 2: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 4: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 6: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + case 8: + hipLaunchKernelGGL((sglang::cross_device_reduce_1stage), dim3(blocks), dim3(threads), 0, stream, + ptrs, fa->sg_, fa->self_sg_, reinterpret_cast(out.data_ptr()), fa->rank_, size); + break; + default: + throw std::runtime_error("world_size must be in (2,4,6,8)"); + } + break; + } +#endif + default: + throw std::runtime_error("deterministic allreduce only supports float32, float16 and bfloat16"); + } +} + +// Deterministic all-reduce for unregistered buffers (ROCm) +void deterministic_all_reduce_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer, torch::Tensor& out) { + const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp)); + auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream(); + + auto input_size = inp.numel() * inp.element_size(); + TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type()); + TORCH_CHECK_EQ(inp.numel(), out.numel()); + TORCH_CHECK(input_size <= reg_buffer.numel() * reg_buffer.element_size(), + "registered buffer is too small to contain the input"); + AT_CUDA_CHECK(hipMemcpyAsync(reg_buffer.data_ptr(), inp.data_ptr(), + input_size, hipMemcpyDeviceToDevice, stream)); + deterministic_all_reduce_reg(_fa, reg_buffer, out); +} diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index 336c3b5a1..245afdb2c 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -51,6 +51,17 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { "()"); m.impl("all_reduce_unreg", torch::kCUDA, &all_reduce_unreg); + // Deterministic all-reduce for ROCm + extern void deterministic_all_reduce_reg(int64_t _fa, torch::Tensor & inp, torch::Tensor & out); + extern void deterministic_all_reduce_unreg( + int64_t _fa, torch::Tensor & inp, torch::Tensor & reg_buffer, torch::Tensor & out); + + m.def("deterministic_all_reduce_reg(int fa, Tensor inp, Tensor! out) -> ()"); + m.impl("deterministic_all_reduce_reg", torch::kCUDA, &deterministic_all_reduce_reg); + + m.def("deterministic_all_reduce_unreg(int fa, Tensor inp, Tensor reg_buffer, Tensor! out) -> ()"); + m.impl("deterministic_all_reduce_unreg", torch::kCUDA, &deterministic_all_reduce_unreg); + m.def("dispose", &dispose); m.def("meta_size", &meta_size); diff --git a/sgl-kernel/python/sgl_kernel/allreduce.py b/sgl-kernel/python/sgl_kernel/allreduce.py index 544fc1d77..b02e8166d 100644 --- a/sgl-kernel/python/sgl_kernel/allreduce.py +++ b/sgl-kernel/python/sgl_kernel/allreduce.py @@ -24,6 +24,18 @@ if torch.version.hip is not None: ) -> None: torch.ops.sgl_kernel.all_reduce_unreg.default(fa, inp, reg_buffer, out) + def deterministic_all_reduce_reg( + fa: int, inp: torch.Tensor, out: torch.Tensor + ) -> None: + torch.ops.sgl_kernel.deterministic_all_reduce_reg.default(fa, inp, out) + + def deterministic_all_reduce_unreg( + fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor + ) -> None: + torch.ops.sgl_kernel.deterministic_all_reduce_unreg.default( + fa, inp, reg_buffer, out + ) + def dispose(fa: int) -> None: torch.ops.sgl_kernel.dispose.default(fa) diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 525b9c0be..2aab5edcf 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -42,6 +42,7 @@ include_dirs = [ sources = [ "csrc/allreduce/custom_all_reduce.hip", + "csrc/allreduce/deterministic_all_reduce.hip", "csrc/allreduce/quick_all_reduce.cu", "csrc/common_extension_rocm.cc", "csrc/elementwise/activation.cu", diff --git a/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py b/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py new file mode 100644 index 000000000..10539fdc3 --- /dev/null +++ b/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py @@ -0,0 +1,270 @@ +""" +Test deterministic custom all-reduce kernel behavior with batch size invariance. + +This test uses the 1-stage all-reduce kernel which is inherently deterministic +due to fixed accumulation ordering (each GPU reads all data from all GPUs and +reduces locally in a fixed order - no atomics, no race conditions). + +Note: This is NOT a reduce-scatter + all-gather (RS+AG) approach. + +This test compares: +1. Deterministic kernel (same batch size) +2. Deterministic kernel (different batch size) + +Usage: + python test_amd_deterministic_custom_allreduce.py +""" + +import multiprocessing as mp +import socket + +import torch +import torch.distributed as dist + + +def get_open_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def worker(world_size, rank, port): + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + ) + + # Try to import and use deterministic kernel + try: + from torch.distributed import new_group + + from sglang.srt.distributed.device_communicators.custom_all_reduce import ( + CustomAllreduce, + ) + + # Create gloo group for custom AR + dist.barrier() + ar_group = new_group(backend="gloo") + dist.barrier() + + custom_ar = CustomAllreduce(group=ar_group, device=device) + + if custom_ar is None or custom_ar.disabled: + if rank == 0: + print("✗ Custom AR not available or disabled") + dist.destroy_process_group() + return + + if not hasattr(custom_ar, "deterministic_all_reduce"): + if rank == 0: + print("✗ Deterministic kernel not available") + dist.destroy_process_group() + return + except Exception as e: + if rank == 0: + print(f"✗ Failed to initialize deterministic kernel: {e}") + import traceback + + traceback.print_exc() + dist.destroy_process_group() + return + + num_trials = 10 + + # Matrix sizes similar to real model layers + # Format: (batch_size, hidden_dim) - typical tensor shape for all-reduce + BS = 50 # max batch_size (1..BS) + hidden_dim = 16384 # hidden dimension / intermediate dimension + + # Different seed per rank - each GPU has DIFFERENT input + torch.manual_seed(42 + rank) + + # Create fixed inputs for all trials + # Single request: (hidden_dim,) + base_input = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device) + base_input_rand = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device) + + # Check if inputs fit in buffer + # Buffer size is max_size bytes, input size is numel * element_size bytes + input_size_bytes = base_input.numel() * base_input.element_size() + if input_size_bytes > custom_ar.max_size and rank == 0: + print( + f"Warning: Input size ({input_size_bytes/(1024*1024):.1f} MB) exceeds buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)" + ) + print(" Using unregistered mode (will copy to buffer)") + + dist.barrier() + + # ========================================================================= + # TEST 1: Deterministic kernel (same batch size) - should be DETERMINISTIC + # ========================================================================= + if rank == 0: + print(f"\n{'='*70}") + print("TEST 1: Deterministic kernel (same batch size)") + print(f"{'='*70}") + dist.barrier() + + results_allreduce_only = [] + for trial in range(num_trials): + # Clone the same input + inp = base_input.clone() + + # Use deterministic kernel + # Check if input fits in buffer, use registered mode if too large + input_size_bytes = inp.numel() * inp.element_size() + use_registered = input_size_bytes > custom_ar.max_size + + if use_registered: + # For large inputs, register buffer first + custom_ar.register_buffer(inp) + result = custom_ar.deterministic_all_reduce(inp, registered=True) + else: + # For smaller inputs, use unregistered mode (copies to internal buffer) + result = custom_ar.deterministic_all_reduce(inp, registered=False) + torch.cuda.synchronize() + + # Store checksum + checksum = result.view(-1).sum().item() + first_vals = result.view(-1)[:5].clone() + results_allreduce_only.append((checksum, first_vals)) + + if rank == 0: + print( + f" Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" + ) + + # Check determinism + if rank == 0: + ref_sum, ref_vals = results_allreduce_only[0] + all_match = True + for i, (s, vals) in enumerate(results_allreduce_only[1:], 1): + if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3): + all_match = False + print(f" Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}") + + if all_match: + print(" ✓ DETERMINISTIC KERNEL (fixed BS): DETERMINISTIC (as expected)") + else: + print( + " ✗ DETERMINISTIC KERNEL (fixed BS): NON-DETERMINISTIC (unexpected!)" + ) + + dist.barrier() + + # ========================================================================= + # TEST 2: Deterministic kernel (different batch size) - should be DETERMINISTIC + # [a], [a, x], [a, x, x], ... + # ========================================================================= + if rank == 0: + print(f"\n{'='*70}") + print("TEST 2: Deterministic kernel (different batch size)") + print("Batches: [a], [a,x], [a,x,x], ...") + print(f"{'='*70}") + dist.barrier() + + results_allreduce_only = {trial: [] for trial in range(num_trials)} + for trial in range(num_trials): + for bs in range(1, BS + 1): + # Construct batch: (batch_size, hidden_dim) + # First element is base_input, rest are base_input_rand + batch = torch.stack([base_input] + [base_input_rand] * (bs - 1), dim=0) + # Shape: (bs, hidden_dim) + + # Flatten for all-reduce: (bs * hidden_dim,) + batch_flat = batch.view(-1) + + # Use deterministic kernel + # Check if input fits in buffer, use registered mode if too large + input_size_bytes = batch_flat.numel() * batch_flat.element_size() + use_registered = input_size_bytes > custom_ar.max_size + + if use_registered: + # For large inputs, register buffer first + custom_ar.register_buffer(batch_flat) + result_flat = custom_ar.deterministic_all_reduce( + batch_flat, registered=True + ) + else: + # For smaller inputs, use unregistered mode + result_flat = custom_ar.deterministic_all_reduce( + batch_flat, registered=False + ) + torch.cuda.synchronize() + + # Reshape back to (bs, hidden_dim) + batch_out = result_flat.view(bs, hidden_dim) + + # Only compare output corresponding to first request + out_first_req = batch_out[0].clone() + checksum = out_first_req.sum().item() + first_vals = out_first_req[:5].clone() + results_allreduce_only[trial].append((bs, checksum, first_vals)) + + if rank == 0: + print( + f" Batch size {bs:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" + ) + + # Check determinism + if rank == 0: + for trial in range(num_trials): + results = results_allreduce_only[trial] + + _, ref_sum, ref_vals = results[0] + all_match = True + for _, s, vals in results[1:]: + if abs(ref_sum - s) > 1e-3 or not torch.allclose( + ref_vals, vals, rtol=1e-3 + ): + all_match = False + + if all_match: + print(" ✓ DETERMINISTIC KERNEL (variant BS): DETERMINISTIC") + else: + print(" ✗ DETERMINISTIC KERNEL (variant BS): NON-DETERMINISTIC") + + dist.barrier() + + dist.destroy_process_group() + + +def main(): + world_size = 8 + available_gpus = torch.cuda.device_count() + + print("=" * 70) + print("Deterministic Kernel All-Reduce Determinism Test") + print("=" * 70) + print(f"Available GPUs: {available_gpus}") + print(f"Using world_size: {world_size}") + + if available_gpus < world_size: + print( + f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead" + ) + world_size = available_gpus + + if world_size < 2: + print("ERROR: Need at least 2 GPUs for this test") + return + + mp.set_start_method("spawn", force=True) + port = get_open_port() + + procs = [] + for rank in range(world_size): + p = mp.Process(target=worker, args=(world_size, rank, port)) + p.start() + procs.append(p) + + for p in procs: + p.join() + + +if __name__ == "__main__": + main() diff --git a/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py b/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py new file mode 100644 index 000000000..dc79d7455 --- /dev/null +++ b/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py @@ -0,0 +1,198 @@ +""" +Test to confirm non-determinism of default NCCL all-reduce with batch size invariance. + +This test uses the default torch.distributed.all_reduce (NCCL) which can be +NON-DETERMINISTIC due to tree-based reduction algorithms that don't guarantee +fixed accumulation order for bfloat16/float16. + +This test compares: +1. Default all-reduce (same batch size) - should be DETERMINISTIC +2. Default all-reduce (different batch size) - typically NON-DETERMINISTIC for bfloat16 + +Usage: + python test_ar.py +""" + +import multiprocessing as mp +import socket + +import torch +import torch.distributed as dist + + +def get_open_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def worker(world_size, rank, port): + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + ) + + num_trials = 10 + + # Matrix sizes similar to real model layers + # Format: (batch_size, hidden_dim) - typical tensor shape for all-reduce + BS = 50 # max batch_size (1..BS) + hidden_dim = 16384 # hidden dimension / intermediate dimension + + # Different seed per rank - each GPU has DIFFERENT input + torch.manual_seed(42 + rank) + + # Create fixed inputs for all trials + # Single request: (hidden_dim,) + base_input = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device) + base_input_rand = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device) + + dist.barrier() + + # ========================================================================= + # TEST 1: Default all-reduce (same batch size) - should be DETERMINISTIC + # ========================================================================= + if rank == 0: + print(f"\n{'='*70}") + print("TEST 1: Default NCCL all_reduce (same batch size)") + print(f"{'='*70}") + dist.barrier() + + results_allreduce_only = [] + for trial in range(num_trials): + # Clone the same input + inp = base_input.clone() + + # Use default NCCL all-reduce + dist.all_reduce(inp) + torch.cuda.synchronize() + + # Store checksum + checksum = inp.view(-1).sum().item() + first_vals = inp.view(-1)[:5].clone() + results_allreduce_only.append((checksum, first_vals)) + + if rank == 0: + print( + f" Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" + ) + + # Check determinism + if rank == 0: + ref_sum, ref_vals = results_allreduce_only[0] + all_match = True + for i, (s, vals) in enumerate(results_allreduce_only[1:], 1): + if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3): + all_match = False + print(f" Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}") + + if all_match: + print(" ✓ DEFAULT ALL_REDUCE (fixed BS): DETERMINISTIC (as expected)") + else: + print(" ✗ DEFAULT ALL_REDUCE (fixed BS): NON-DETERMINISTIC (unexpected!)") + + dist.barrier() + + # ========================================================================= + # TEST 2: Default all-reduce (different batch size) - typically NON-DETERMINISTIC + # [a], [a, x], [a, x, x], ... + # ========================================================================= + if rank == 0: + print(f"\n{'='*70}") + print("TEST 2: Default NCCL all_reduce (different batch size)") + print("Batches: [a], [a,x], [a,x,x], ...") + print(f"{'='*70}") + dist.barrier() + + results_allreduce_only = {trial: [] for trial in range(num_trials)} + for trial in range(num_trials): + for bs in range(1, BS + 1): + # Construct batch: (batch_size, hidden_dim) + # First element is base_input, rest are base_input_rand + batch = torch.stack([base_input] + [base_input_rand] * (bs - 1), dim=0) + # Shape: (bs, hidden_dim) + + # Flatten for all-reduce: (bs * hidden_dim,) + batch_flat = batch.view(-1) + + # Use default NCCL all-reduce + dist.all_reduce(batch_flat) + torch.cuda.synchronize() + + # Reshape back to (bs, hidden_dim) + batch_out = batch_flat.view(bs, hidden_dim) + + # Only compare output corresponding to first request + out_first_req = batch_out[0].clone() + checksum = out_first_req.sum().item() + first_vals = out_first_req[:5].clone() + results_allreduce_only[trial].append((bs, checksum, first_vals)) + + if rank == 0: + print( + f" Batch size {bs:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" + ) + + # Check determinism + if rank == 0: + for trial in range(num_trials): + results = results_allreduce_only[trial] + + _, ref_sum, ref_vals = results[0] + all_match = True + for _, s, vals in results[1:]: + if abs(ref_sum - s) > 1e-3 or not torch.allclose( + ref_vals, vals, rtol=1e-3 + ): + all_match = False + + if all_match: + print(" ✓ DEFAULT ALL_REDUCE (variant BS): DETERMINISTIC") + else: + print(" ✗ DEFAULT ALL_REDUCE (variant BS): NON-DETERMINISTIC") + + dist.barrier() + + dist.destroy_process_group() + + +def main(): + world_size = 8 + available_gpus = torch.cuda.device_count() + + print("=" * 70) + print("Default NCCL All-Reduce Determinism Test") + print("=" * 70) + print(f"Available GPUs: {available_gpus}") + print(f"Using world_size: {world_size}") + + if available_gpus < world_size: + print( + f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead" + ) + world_size = available_gpus + + if world_size < 2: + print("ERROR: Need at least 2 GPUs for this test") + return + + mp.set_start_method("spawn", force=True) + port = get_open_port() + + procs = [] + for rank in range(world_size): + p = mp.Process(target=worker, args=(world_size, rank, port)) + p.start() + procs.append(p) + + for p in procs: + p.join() + + +if __name__ == "__main__": + main()