Refactor: decouple segment tracking from comm registration (#21392)

Signed-off-by: wangfakang <fakangwang@gmail.com>

benchmark/bench_pynccl_allocator/bench_segment_tracking.py

@@ -0,0 +1,210 @@
+"""
+Benchmark for comparing CPU overhead of segment tracking methods:
+1. nccl_allocator_register_segments_with_comm() - C++ registration with index tracking
+2. torch.cuda.memory.memory_snapshot() - PyTorch memory snapshot
+
+Usage:
+    python benchmark/bench_pynccl_allocator/bench_segment_tracking.py --num-segments 50 --num-iters 1000
+"""
+
+import argparse
+import time
+import warnings
+from typing import List
+
+import torch
+
+warnings.filterwarnings("ignore")
+
+
+def setup_segments(num_segments: int, segment_size: int = 1024 * 1024):
+    """
+    Allocate a specified number of segments using the NCCL allocator.
+    """
+    import os
+
+    import torch.distributed as dist
+
+    from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+        get_nccl_mem_pool,
+    )
+
+    # Initialize distributed if not already done
+    if not dist.is_initialized():
+        os.environ.setdefault("MASTER_ADDR", "localhost")
+        os.environ.setdefault("MASTER_PORT", "29500")
+        dist.init_process_group(
+            backend="nccl",
+            rank=0,
+            world_size=1,
+            device_id=torch.device(f"cuda:{torch.cuda.current_device()}"),
+        )
+
+    mem_pool = get_nccl_mem_pool()
+
+    # Allocate segments in the pool
+    tensors: List[torch.Tensor] = []
+    with torch.cuda.use_mem_pool(mem_pool):
+        for _ in range(num_segments):
+            t = torch.empty(segment_size, dtype=torch.uint8, device="cuda")
+            tensors.append(t)
+
+    # Keep tensors alive by returning them (caller should hold reference)
+    return tensors, mem_pool
+
+
+def bench_register_segments_with_comm(
+    nccl_lib, comm_ptr: int, num_iters: int = 10000
+) -> float:
+    """
+    Benchmark nccl_allocator_register_segments_with_comm() function.
+
+    Args:
+        nccl_lib: The loaded NCCL allocator library
+        comm_ptr: The communicator pointer value
+        num_iters: Number of iterations
+
+    Returns:
+        Average time per call in microseconds.
+    """
+    import ctypes
+
+    # Setup the C function signature
+    register_func = nccl_lib.nccl_allocator_register_segments_with_comm
+    register_func.restype = ctypes.c_int
+    register_func.argtypes = [ctypes.c_uint64]
+
+    # Warmup
+    for _ in range(100):
+        register_func(comm_ptr)
+
+    # Benchmark
+    start = time.perf_counter()
+    for _ in range(num_iters):
+        register_func(comm_ptr)
+    end = time.perf_counter()
+
+    avg_us = (end - start) / num_iters * 1e6
+    return avg_us
+
+
+def bench_mempool_snapshot(
+    mem_pool: torch.cuda.MemPool, num_iters: int = 10000
+) -> float:
+    """
+    Benchmark torch.cuda.MemPool.snapshot() function.
+
+    Returns:
+        Average time per call in microseconds.
+    """
+    # Warmup
+    for _ in range(100):
+        mem_pool.snapshot()
+
+    # Benchmark
+    start = time.perf_counter()
+    for _ in range(num_iters):
+        mem_pool.snapshot()
+    end = time.perf_counter()
+
+    avg_us = (end - start) / num_iters * 1e6
+    return avg_us
+
+
+def bench_with_various_segment_counts(
+    segment_counts: List[int],
+    num_iters: int = 10000,
+    segment_size: int = 1024 * 1024,  # 1MB per segment
+):
+    """
+    Run benchmarks with various numbers of tracked segments.
+    """
+    print("=" * 80)
+    print("Benchmark: Segment Registration CPU Overhead")
+    print("=" * 80)
+    print(f"Segment size: {segment_size / 1024 / 1024:.2f} MB")
+    print(f"Iterations per measurement: {num_iters}")
+    print()
+    print(
+        f"{'Segments':<12} {'register_segments (µs)':<30} {'snapshot (µs)':<20} {'Speedup':<10}"
+    )
+    print("-" * 80)
+
+    all_tensors = []  # Keep all tensors alive
+    comm_ptr = 0  # Use dummy comm_ptr for benchmarking (no actual NCCL registration)
+
+    for num_segments in segment_counts:
+        # Clean up previous segments
+        all_tensors = []
+
+        # Allocate segments (this initializes _nccl_allocator_lib via get_nccl_mem_pool)
+        tensors, mem_pool = setup_segments(num_segments, segment_size)
+        all_tensors.extend(tensors)
+
+        # Sync to ensure allocations are complete
+        torch.cuda.synchronize()
+
+        # Import _nccl_allocator_lib after setup_segments (ensures library is loaded)
+        from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+            _nccl_allocator_lib,
+        )
+
+        # Run benchmarks
+        time_register = bench_register_segments_with_comm(
+            _nccl_allocator_lib, comm_ptr, num_iters
+        )
+        time_snapshot = bench_mempool_snapshot(mem_pool, num_iters)
+
+        speedup = time_snapshot / time_register if time_register > 0 else float("inf")
+
+        print(
+            f"{num_segments:<12} {time_register:<30.3f} {time_snapshot:<20.3f} {speedup:<10.2f}x"
+        )
+
+    print("-" * 80)
+    print()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Benchmark segment tracking methods in pynccl_allocator"
+    )
+    parser.add_argument(
+        "--num-segments",
+        type=int,
+        nargs="+",
+        default=[10, 50, 100, 200, 500, 1000],
+        help="Number of segments to track (can specify multiple values)",
+    )
+    parser.add_argument(
+        "--num-iters",
+        type=int,
+        default=10000,
+        help="Number of iterations for each measurement",
+    )
+    parser.add_argument(
+        "--segment-size",
+        type=int,
+        default=1024 * 1024,  # 1MB
+        help="Size of each segment in bytes",
+    )
+    args = parser.parse_args()
+
+    # Check CUDA availability
+    if not torch.cuda.is_available():
+        print("Error: CUDA is not available. This benchmark requires a GPU.")
+        return
+
+    # Initialize CUDA context by creating a small tensor
+    _ = torch.zeros(1, device="cuda")
+
+    # Run benchmarks
+    bench_with_various_segment_counts(
+        segment_counts=args.num_segments,
+        num_iters=args.num_iters,
+        segment_size=args.segment_size,
+    )
+
+
+if __name__ == "__main__":
+    main()