[JIT Kernel] Multi-GPU test/bench framework for custom all-reduce + TP QKNorm (#26706)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai>

python/sglang/jit_kernel/benchmark/marker.py

@@ -1,3 +1,4 @@
+import contextlib
 import inspect
 import itertools
 import math
@@ -5,6 +6,7 @@ import os
 from typing import (
     Any,
     Callable,
+    ContextManager,
     Dict,
     Generic,
     Iterable,
@@ -93,6 +95,60 @@ def _process_metrics(times: list[float], metrics: tuple[Metric, ...]) -> list[fl
     return results
 
 
+@cache_once
+def _get_l2_cache_size() -> int:
+    device = torch.cuda.current_device()
+    props = torch.cuda.get_device_properties(device)
+    return props.L2_cache_size
+
+
+_L2_SAFE_RATIO = 5
+
+
+def _get_flush_l2_buffer() -> torch.Tensor:
+    """Get a buffer sized to flush the L2 cache when accessed."""
+    device = torch.device(f"cuda:{torch.cuda.current_device()}")
+    l2_size = _get_l2_cache_size()
+    safe_size = int(l2_size * _L2_SAFE_RATIO)
+    return torch.empty(safe_size, device=device, dtype=torch.uint8)
+
+
+def _calculate_rotation_count(nbytes: int, min_rotations: int = 2) -> int:
+    """
+    Adapted from flashinfer benchmark utility:
+    https://github.com/flashinfer-ai/flashinfer/blob/c5a2b06edae4fa2bfd2ae25eed16eb565c70513f/flashinfer/testing/utils.py
+
+    Calculate the number of buffer copies needed to ensure cold L2 cache.
+
+    The function uses conservative thresholds to account for:
+    - LRU eviction being gradual (not all data evicted when capacity exceeded)
+    - Cache associativity effects (some data may persist in non-conflicting sets)
+    - Hardware prefetching behavior
+
+    Returns 1 (no rotation needed) only when tensor size substantially exceeds
+    L2 cache, ensuring cache effects are truly negligible.
+
+    Args:
+        tensors: List of tensors to consider for rotation (must be on GPU).
+        device: Device for L2 cache query (None for current device).
+        min_rotations: Minimum number of rotations when rotation is needed.
+
+    Returns:
+        Number of buffer copies needed (1 means no rotation needed).
+    """
+    l2_size = _get_l2_cache_size()
+    safe_cache_threshold = l2_size * _L2_SAFE_RATIO
+
+    if nbytes <= 0 or nbytes >= safe_cache_threshold:
+        return 1  # No tensors to rotate
+
+    # Conservative formula: ensure between any two uses of the same buffer,
+    # we've accessed enough data to fully flush L2 with margin
+    # Using safe_cache_threshold ensures we account for all cache effects
+    num_rotations = math.ceil(safe_cache_threshold / nbytes) + 1
+    return max(min_rotations, num_rotations)
+
+
 class BenchResult(NamedTuple):
     metrics: Tuple[Metric, ...]
     times: List[float]  # in seconds
@@ -319,6 +375,67 @@ def parametrize(names: str, vals: List[Any], ci_vals: Optional[List[Any]] = None
     return decorator
 
 
+def _do_bench_internal_graph(
+    fn: Callable,
+    replay_iters: int,
+    input_args: Tuple[Any, ...],
+    input_kwargs: Dict[str, Any],
+    graph_clone_args: Iterable[int],
+    graph_clone_kwargs: Iterable[str],
+    graph_context: ContextManager,
+    sync_multigpu_fn: Callable[[], Any],
+) -> List[float]:
+    result: List[float] = []
+    stream = torch.cuda.current_stream()
+    empty_tensor = _get_flush_l2_buffer()
+    # only count the cloned tensors for rotation count
+    nbytes = sum(_get_nbytes_recursive(input_args[i]) for i in graph_clone_args)
+    nbytes += sum(_get_nbytes_recursive(input_kwargs[k]) for k in graph_clone_kwargs)
+    rotate_count = min(_calculate_rotation_count(nbytes), 100)
+    loop_count = math.ceil(100 / rotate_count) * rotate_count
+    input_args_list = [input_args] * rotate_count
+    input_kwargs_list = [input_kwargs] * rotate_count
+    graph_clone_args = set(graph_clone_args)
+    graph_clone_kwargs = set(graph_clone_kwargs)
+
+    graph = torch.cuda.CUDAGraph()
+    # NOTE: we rotate the buffer here to avoid L2 cache effect
+    for i in range(1, rotate_count):
+        input_args_list[i] = tuple(
+            (
+                _clone_recursive(input_args[j])
+                if j in graph_clone_args
+                else input_args[j]
+            )
+            for j in range(len(input_args))
+        )
+        input_kwargs_list[i] = dict(
+            (k, (_clone_recursive(v) if k in graph_clone_kwargs else v))
+            for k, v in input_kwargs.items()
+        )
+    with graph_context:
+        with torch.cuda.graph(graph, stream=stream):
+            for i in range(loop_count):
+                args = input_args_list[i % rotate_count]
+                kwargs = input_kwargs_list[i % rotate_count]
+                fn(*args, **kwargs)
+
+    # warm up the graph once
+    graph.replay()
+    # then replay the graph and measure the time
+    tic = torch.cuda.Event(enable_timing=True)
+    toc = torch.cuda.Event(enable_timing=True)
+    for _ in range(max(replay_iters // loop_count, 10)):
+        empty_tensor.zero_()  # cold the L2 cache
+        sync_multigpu_fn()  # sync GPU before each iteration for precise timing
+        tic.record(stream)
+        graph.replay()
+        toc.record(stream)
+        stream.synchronize()
+        result.append(tic.elapsed_time(toc) / loop_count)
+    return result
+
+
 def do_bench(
     fn: Callable,
     *,
@@ -338,9 +455,13 @@ def do_bench(
     memory_output: Iterable[Any] | Literal["out"] | None = "out",
     extra_memory_args: Iterable[Any] | None = None,
     extra_memory_footprint: int = 0,
+    graph_context_fn: Optional[Callable[[], ContextManager]] = None,
+    sync_multigpu_fn: Optional[Callable[[], Any]] = None,
 ) -> BenchResult:
     """
     Benchmark a function using CUDA graph or naive loop.
+    Adapted from flashinfer benchmark utility:
+    https://github.com/flashinfer-ai/flashinfer/blob/c5a2b06edae4fa2bfd2ae25eed16eb565c70513f/flashinfer/testing/utils.py
 
     :param fn: Function to benchmark
     :param input_args: Positional arguments to pass to the function
@@ -361,6 +482,10 @@ def do_bench(
     :param extra_memory_args: Additional arguments to consider for memory footprint calculation.
     :param extra_memory_footprint: Additional memory footprint to consider.
                                    This is typically used when the load/store bytes is dynamic.
+    :param graph_context_fn: A callable returning a context manager that wraps the cuda graph capture.
+    :param sync_multigpu_fn: A callable to synchronize multiple GPUs before each iteration. For precise
+                             benchmark number in multi-GPU benchmark, it should be some synchronization
+                             primitive on GPU side (not on CPU side).
     """
     # first warmup the function
     device_id = torch.cuda.current_device()
@@ -368,17 +493,14 @@ def do_bench(
         stream = _get_benchmark_stream(device_id)
     old_current_stream = torch.cuda.current_stream(device_id)
     result: List[float] = []
+    sync_multigpu_fn = sync_multigpu_fn or (lambda: None)
     with torch.cuda.device(device_id), torch.cuda.stream(stream):
         stream.wait_stream(old_current_stream)
+        sync_multigpu_fn()
         for _ in range(warmup_iters):
             fn(*input_args, **input_kwargs)
         if use_cuda_graph:
             # NOTE: by default, reduce all the CPU-side overhead
-            rep_count = 4
-            loop_iters = 100
-            graph = torch.cuda.CUDAGraph()
-            input_args_list = [input_args] * rep_count
-            input_kwargs_list = [input_kwargs] * rep_count
             if graph_clone_args == "all":
                 graph_clone_args = range(len(input_args))
             elif graph_clone_args is None:
@@ -387,44 +509,29 @@ def do_bench(
                 graph_clone_kwargs = input_kwargs.keys()
             elif graph_clone_kwargs is None:
                 graph_clone_kwargs = []
-            graph_clone_args = set(graph_clone_args)
-            graph_clone_kwargs = set(graph_clone_kwargs)
-            # NOTE: we rotate the buffer here to avoid L2 cache effect
-            for i in range(1, rep_count):
-                input_args_list[i] = tuple(
-                    (
-                        _clone_recursive(input_args[j])
-                        if j in graph_clone_args
-                        else input_args[j]
-                    )
-                    for j in range(len(input_args))
-                )
-                input_kwargs_list[i] = dict(
-                    (k, (_clone_recursive(v) if k in graph_clone_kwargs else v))
-                    for k, v in input_kwargs.items()
-                )
-            with torch.cuda.graph(graph, stream=stream):
-                for _ in range(loop_iters // rep_count):
-                    for args, kwargs in zip(input_args_list, input_kwargs_list):
-                        fn(*args, **kwargs)
-            # warm up the graph
-            graph.replay()
-            # then replay the graph and measure the time
-            tic = torch.cuda.Event(enable_timing=True)
-            toc = torch.cuda.Event(enable_timing=True)
-            for _ in range(max(replay_iters // loop_iters, 10)):
-                tic.record(stream)
-                graph.replay()
-                toc.record(stream)
-                stream.synchronize()
-                result.append(tic.elapsed_time(toc) / loop_iters)
+            graph_context = (
+                graph_context_fn()
+                if graph_context_fn is not None
+                else contextlib.nullcontext()
+            )
+            result = _do_bench_internal_graph(
+                fn,
+                replay_iters,
+                input_args,
+                input_kwargs,
+                graph_clone_args,
+                graph_clone_kwargs,
+                graph_context,
+                sync_multigpu_fn,
+            )
         else:
             # NOTE: no cuda graph, naive loop
-            empty_tensor = torch.empty(64 * 1024 * 1024, device=f"cuda:{device_id}")
             tic = torch.cuda.Event(enable_timing=True)
             toc = torch.cuda.Event(enable_timing=True)
+            empty_tensor = _get_flush_l2_buffer()
             for _ in range(max(replay_iters, 10)):
                 empty_tensor.zero_()  # cold the L2 cache
+                sync_multigpu_fn()
                 tic.record(stream)
                 fn(*input_args, **input_kwargs)
                 toc.record(stream)

python/sglang/jit_kernel/benchmark/utils.py

@@ -1,12 +1,64 @@
 """Common utilities for jit_kernel benchmark files."""
 
-from typing import Callable, List, Sequence, Tuple
+from typing import Callable, List, Optional, Sequence, Tuple
 
 import torch
 import triton.testing
 
+from sglang.jit_kernel.mp import multigpu_launch
 from sglang.utils import is_in_ci
 
+
+def multigpu_bench_main(
+    name: str,
+    file: str,
+    num_gpus: Sequence[int],
+    main_fn: Callable[[], None],
+    *,
+    pre_launch_fn: Optional[Callable[[List[int]], None]] = None,
+    timeout: Optional[int] = None,
+) -> None:
+    """cudalib-style multi-GPU benchmark entry point.
+
+    Drop this at the bottom of a benchmark file::
+
+        multigpu_bench_main(
+            name=__name__,
+            file=__file__,
+            num_gpus=range(2, 9),
+            main_fn=benchmark.run,
+        )
+
+    Mirrors :func:`multigpu_pytest_main` but invokes a caller-supplied function
+    instead of pytest. ``main_fn`` is expected to return ``None`` on success;
+    any exception propagates as a non-zero exit. Pass ``--num-gpu 2,4`` on the
+    command line to override ``num_gpus``.
+
+    ``pre_launch_fn`` (kw-only) runs once in the outer process before any
+    torchrun child starts, receiving the runnable world sizes. Use it for
+    parallel JIT precompilation so torchrun children hit a warm disk cache.
+
+    ``timeout`` (kw-only, seconds) bounds each per-world-size torchrun
+    invocation. Defaults to ``None`` (wait indefinitely) since benchmark sweeps
+    can legitimately run long; set it to fail fast on a hung worker.
+    """
+
+    def inner() -> int:
+        main_fn()
+        return 0
+
+    return multigpu_launch(
+        name,
+        file,
+        num_gpus,
+        env_key="_IS_BENCH_MULTIGPU_SGLANG_JIT_KERNEL",
+        inner=inner,
+        kind="benchmark",
+        pre_launch_fn=pre_launch_fn,
+        timeout=timeout,
+    )
+
+
 # Common constants
 DEFAULT_DTYPE = torch.bfloat16
 DEFAULT_DEVICE = "cuda"

python/sglang/jit_kernel/mp.py

@@ -0,0 +1,214 @@
+"""Multi-process / multi-GPU launching utilities (torchrun-based).
+
+Shared `multigpu_launch` helper that both `sglang.jit_kernel.tests.utils` and
+`sglang.jit_kernel.benchmark.utils` build their domain-specific entry points on
+top of (`multigpu_pytest_main`, `multigpu_bench_main`).
+
+When a script that calls one of those wrappers is run with plain `python`, the
+launcher relaunches the same file under `torchrun` once for each `N` in
+`num_gpus`. When the inner workers run (identified by an env_key being set),
+the same launcher calls `inner()` on every rank, silences stdout on non-zero
+ranks, and exits with its return code.
+"""
+
+from __future__ import annotations
+
+import atexit
+import logging
+import os
+import signal
+import subprocess
+import sys
+from typing import Any, Callable, List, NoReturn, Optional, Sequence
+
+import psutil
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def register_comm_cleanup(comm: Any) -> None:
+    """Register an idempotent shutdown for a custom-AR communicator."""
+
+    def _safe_close() -> None:
+        try:
+            comm.close()
+        except Exception:
+            pass
+        # Disable both class flavors' early-out paths in __del__/close.
+        try:
+            comm.disabled = True
+        except Exception:
+            pass
+        # CustomAllReduceV2: drop ``obj`` so close() short-circuits next time.
+        try:
+            delattr(comm, "obj")
+        except Exception:
+            pass
+        # CustomAllreduce: zero ``_ptr`` so close() short-circuits next time.
+        try:
+            comm._ptr = 0
+        except Exception:
+            pass
+
+    atexit.register(_safe_close)
+
+
+def _kill_pgroup(pgid: int) -> None:
+    try:
+        os.killpg(pgid, signal.SIGKILL)
+    except ProcessLookupError:
+        pass
+
+
+def _kill_descendants(pid: int) -> None:
+    """Snapshot every descendant of `pid` *now* and SIGKILL them all.
+
+    Must be called BEFORE the direct child (torchrun) dies -- once it does,
+    its workers get reparented to init and we lose them via the process tree.
+    """
+    try:
+        root = psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        return
+    descendants = root.children(recursive=True)
+    for proc in descendants:
+        try:
+            proc.kill()
+        except psutil.Error:
+            # NoSuchProcess (already gone) or AccessDenied -- nothing to do.
+            pass
+    psutil.wait_procs(descendants, timeout=5)
+
+
+def _extract_num_gpus_override(
+    argv: list[str],
+) -> tuple[list[int] | None, list[str]]:
+    """Pop `--num-gpu(s)` flags out of `argv` and return them separately.
+
+    Accepts `--num-gpu N`, `--num-gpu=N`, `--num-gpus ...`, and comma-separated
+    lists like `--num-gpu 2,4,8`. May be repeated.
+    """
+    override: list[int] = []
+    remaining: list[str] = []
+    i = 0
+    while i < len(argv):
+        a = argv[i]
+        if a in ("--num-gpu", "--num-gpus"):
+            if i + 1 >= len(argv):
+                raise ValueError(f"missing value for {a} (expected e.g. `{a} 2,4`)")
+            override.extend(int(x) for x in argv[i + 1].split(","))
+            i += 2
+        elif a.startswith("--num-gpu=") or a.startswith("--num-gpus="):
+            _, val = a.split("=", 1)
+            override.extend(int(x) for x in val.split(","))
+            i += 1
+        else:
+            remaining.append(a)
+            i += 1
+    return (override if override else None), remaining
+
+
+def multigpu_launch(
+    name: str,
+    file: str,
+    num_gpus: Sequence[int],
+    env_key: str,
+    inner: Callable[[], int],
+    kind: str,
+    pre_launch_fn: Optional[Callable[[List[int]], None]] = None,
+    timeout: Optional[int] = None,
+) -> NoReturn | None:
+    """Shared torchrun-based launcher.
+
+    See module docstring. `name` is the caller's `__name__`; `file` is its
+    `__file__`. `env_key` is a unique string per kind (test/benchmark) used to
+    detect the inside-torchrun state. `inner` returns an exit code.
+
+    `pre_launch_fn`, if given, runs once in the outer process *before* any
+    torchrun child is spawned. It receives the list of world sizes that will
+    actually be launched (already filtered against the host's GPU count and
+    any ``--num-gpu`` override). Use it for parallel JIT precompilation so the
+    on-disk kernel cache is warm by the time the torchrun children import
+    their kernels.
+
+    `timeout`, if given, bounds each per-world-size torchrun invocation (in
+    seconds). On expiry the child's whole process group is killed and the
+    launcher exits non-zero. `None` (the default) waits indefinitely.
+    """
+    pid_key = env_key + "_PID"
+    if env_key in os.environ:
+        assert pid_key in os.environ
+        if name != "__main__":
+            return
+        rank = int(os.environ["LOCAL_RANK"])
+        if rank != 0:
+            sys.stdout = open(os.devnull, "w")
+        torch.cuda.set_device(rank)
+        return sys.exit(inner())
+    assert pid_key not in os.environ
+    if name != "__main__":
+        return logger.warning(
+            f"{file} can not directly run with `pytest`. "
+            "Use `python` to invoke it, which will internally relaunch it "
+            "under torchrun for each requested number of GPUs."
+        )
+    num_devices = torch.cuda.device_count()
+    override, forwarded_args = _extract_num_gpus_override(sys.argv[1:])
+    if override is not None:
+        logger.info(f"--num-gpu override: running only with {override}")
+        num_gpus = override
+        for N in num_gpus:
+            if N <= 1 or N > num_devices:
+                raise ValueError(
+                    f"Invalid number of GPUs requested: {N} "
+                    f"(available: {num_devices})"
+                )
+    os.environ[env_key] = "1"
+    os.environ[pid_key] = str(os.getpid())
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
+    os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo")  # single-machine setup
+    signal.signal(signal.SIGINT, signal.default_int_handler)
+    runnable: List[int] = []
+    for N in sorted(num_gpus):
+        assert N > 1
+        if N > num_devices:
+            logger.warning(f"Skipping {kind} with {N} GPUs ({num_devices} available)")
+            continue
+        runnable.append(N)
+    if pre_launch_fn is not None and runnable:
+        logger.info(f"Running pre-launch hook for world sizes {runnable}")
+        pre_launch_fn(runnable)
+    for N in runnable:
+        logger.info(f"Running {kind} with {N} GPUs")
+        cmd = [
+            "torchrun",
+            "--nproc_per_node",
+            str(N),
+            "--local-addr",
+            "127.0.0.1",
+            file,
+        ]
+        cmd += forwarded_args
+        proc = subprocess.Popen(cmd, start_new_session=True)
+        pgid = proc.pid
+        returncode = -1
+        timed_out = False
+        try:
+            returncode = proc.wait(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            timed_out = True
+        finally:
+            _kill_descendants(os.getpid())
+            _kill_pgroup(pgid)
+            try:
+                proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                pass
+        if timed_out:
+            logger.error(f"{kind} (nproc={N}) timed out after {timeout} seconds")
+            sys.exit(1)
+        if returncode != 0:
+            logger.error(f"{kind} failed with {N} GPUs (exit {returncode})")
+            sys.exit(returncode)
+    logger.info(f"All {kind}s passed")

python/sglang/jit_kernel/tests/utils.py

@@ -1,49 +1,57 @@
-import os
-import subprocess
 import sys
-from typing import Callable
+from typing import Callable, List, Optional, Sequence
 
 import pytest
 
+from sglang.jit_kernel.mp import multigpu_launch
 
-def multiprocess_test(file: str, nproc: int, timeout: int = 240) -> None:
-    """Launch this script as a torchrun worker and assert success.
 
-    The default budget covers the cold-cache first invocation, where the
-    worker pays the full triton + cutlass JIT compile cost (60-180s observed
-    on H200). The previous 90s default tripped intermittently on the first
-    parametrisation of `test_tp_qknorm` (seen on `main` runs too, not only
-    on fresh-venv PRs); subsequent parametrisations finished in ~60s once
-    the JIT cache was warm.
+def multigpu_pytest_main(
+    name: str,
+    file: str,
+    num_gpus: Sequence[int],
+    *,
+    pre_launch_fn: Optional[Callable[[List[int]], None]] = None,
+    timeout: Optional[int] = 600,
+) -> None:
+    """cudalib-style multi-GPU pytest entry point.
+
+    Drop this at the bottom of a test file::
+
+        multigpu_pytest_main(__name__, __file__, num_gpus=range(2, 9))
+
+    When the file is run with ``python <file>``, it relaunches itself under
+    ``torchrun --nproc_per_node=N <file>`` for each N in ``num_gpus``. Inside
+    each worker, ``pytest.main([file, ...forwarded_args])`` runs the collected
+    tests. Pass ``--num-gpu 2,4`` on the command line to override ``num_gpus``.
+
+    ``pre_launch_fn`` (kw-only) runs once in the outer process before any
+    torchrun child starts, receiving the runnable world sizes. Use it for
+    parallel JIT precompilation so torchrun children hit a warm disk cache
+    instead of compiling kernels on first call.
+
+    ``timeout`` (kw-only, seconds) bounds each per-world-size torchrun
+    invocation. The default budget covers the cold-cache first invocation
+    (the worker pays the full triton + cutlass JIT compile cost, 60-180s
+    observed on H200) plus the nightly full sweep, which runs every size x
+    dtype x algo x graph-mode parametrisation rather than the reduced in-CI
+    range. A worker that exceeds the budget is killed and the run fails. Pass
+    ``None`` to wait indefinitely.
     """
-    cmd = [
-        "torchrun",
-        f"--nproc_per_node={nproc}",
+
+    def inner() -> int:
+        # CI's run_unittest_files invokes `python3 <file> -f` (legacy
+        # unittest failfast). Translate to pytest's `-x` so it survives.
+        pytest_args = ["-x" if a == "-f" else a for a in sys.argv[1:]]
+        return pytest.main([file] + pytest_args)
+
+    return multigpu_launch(
+        name,
         file,
-    ]
-    try:
-        result = subprocess.run(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            text=True,
-            timeout=timeout,
-        )
-    except subprocess.TimeoutExpired as e:
-        raise RuntimeError(
-            f"torchrun (nproc={nproc}) timed out after {timeout} seconds\n"
-            f"{e.stdout}"
-        ) from e
-
-    assert result.returncode == 0, (
-        f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n"
-        f"{result.stdout}"
+        num_gpus,
+        env_key="_IS_TEST_MULTIGPU_SGLANG_JIT_KERNEL",
+        inner=inner,
+        kind="test",
+        pre_launch_fn=pre_launch_fn,
+        timeout=timeout,
     )
-
-
-def multiprocess_main(file: str, main: Callable[[], None]) -> None:
-    """Helper to run a function in a multiprocess torchrun context."""
-    if "LOCAL_RANK" in os.environ:
-        main()
-    else:
-        sys.exit(pytest.main([file, "-v", "-s"]))