sgl-project--sglang/python/sglang/jit_kernel/benchmark/marker.py

import contextlib
import inspect
import itertools
import math
import os
from typing import (
    Any,
    Callable,
    ContextManager,
    Dict,
    Generic,
    Iterable,
    List,
    Literal,
    NamedTuple,
    Optional,
    Tuple,
    TypeAlias,
    TypeVar,
)

import torch

from sglang.jit_kernel.utils import cache_once
from sglang.utils import is_in_ci

F = TypeVar("F", bound=Callable[..., "BenchResult"])
Metric: TypeAlias = "float | Literal['avg']"
BENCH_CONFIG: TypeAlias = "List[Tuple[Tuple[str, ...], List[Tuple[Any, ...]]]]"
UNIT_SCALE = {"us": 1e-6, "ms": 1e-3, "s": 1.0}
TYPE_LIST = (bool, int, float, str, torch.dtype, torch.device, None.__class__)
DISABLE_LOG_BANDWIDTH = os.environ.get("SGLANG_KERNEL_DISABLE_LOG_BANDWIDTH") == "1"


__all__ = [
    "BenchResult",
    "BenchSkip",
    "Benchmark",
    "benchmark",
    "parametrize",
    "do_bench",
    "skip",
]


class BenchSkip(Exception):
    pass


def skip(reason: str):
    raise BenchSkip(reason)


@cache_once
def _get_benchmark_stream(device_id: int) -> torch.cuda.Stream:
    return torch.cuda.Stream(device=device_id)


def _clone_recursive(in_: Any) -> Any:
    if isinstance(in_, torch.Tensor):
        return in_.clone()
    elif isinstance(in_, (list, tuple)):
        return type(in_)(_clone_recursive(x) for x in in_)
    elif isinstance(in_, dict):
        return {k: _clone_recursive(v) for k, v in in_.items()}
    elif isinstance(in_, TYPE_LIST):
        return in_
    # NOTE: avoid silent error
    raise ValueError(f"unsupported type: {type(in_)}")


def _get_nbytes_recursive(in_: Any) -> int:
    if isinstance(in_, torch.Tensor):
        return in_.nbytes
    elif isinstance(in_, (list, tuple)):
        return sum(_get_nbytes_recursive(x) for x in in_)
    elif isinstance(in_, dict):
        return sum(_get_nbytes_recursive(v) for v in in_.values())
    elif isinstance(in_, TYPE_LIST):
        return 0
    # NOTE: avoid silent error
    raise ValueError(f"unsupported type: {type(in_)}")


def _process_metrics(times: list[float], metrics: tuple[Metric, ...]) -> list[float]:
    results: list[float] = []
    times = sorted(x / 1000 for x in times)  # convert to seconds and sort
    for metric in metrics:
        if metric == "avg":
            results.append(sum(times) / len(times))
        else:
            assert 0 <= metric <= 1, f"invalid metric: {metric}"
            which = min(int(len(times) * metric), len(times) - 1)
            results.append(times[which])
    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
    memory_footprint: Optional[int]


class Table:
    """Aligned text table with `|` section separators and `=`/`-` rules."""

    SEP = " | "

    def __init__(self) -> None:
        self._headers: List[str] = []
        self._mins: List[int] = []
        self._pads: List[int] = []
        self._aligns: List[str] = []
        self._seps: set = set()
        self._rows: List[List[str]] = []

    @staticmethod
    def format_latency(r: float) -> str:
        if math.isnan(r):
            return "N/A"
        length = len(str(int(r)))
        if length < 5:
            return f"{r:.4f}"
        # decrease number of the digits
        digits = max(0, 4 - (length - 5))
        return f"{r:.{digits}f}"

    @staticmethod
    def format_bandwidth(b: float) -> str:
        if math.isnan(b):
            return "N/A"
        return f"{b:.2f}"

    def col(
        self,
        header: str = "",
        *,
        min_width: int = 10,
        pad: int = 2,
        align: str = ">",
    ) -> None:
        self._headers.append(header)
        self._mins.append(min_width)
        self._pads.append(pad)
        self._aligns.append(align)

    def sep(self) -> None:
        self._seps.add(len(self._headers))

    def row(self, *cells: Any) -> None:
        assert len(cells) == len(self._headers)
        self._rows.append([str(c) for c in cells])

    def print(self) -> None:
        widths = [
            max(max(len(c) + p for c in [h, *(r[i] for r in self._rows)]), mw)
            for i, (h, mw, p) in enumerate(zip(self._headers, self._mins, self._pads))
        ]
        total = sum(widths) + len(self.SEP) * len(self._seps)

        def fmt(cells: List[str]) -> str:
            parts: List[str] = []
            for i, (cell, w, a) in enumerate(zip(cells, widths, self._aligns)):
                if i in self._seps:
                    parts.append(self.SEP)
                parts.append(f"{cell:{a}{w}}")
            return "".join(parts)

        print("=" * total)
        print(fmt(self._headers))
        print("-" * total)
        for r in self._rows:
            print(fmt(r))
        print("=" * total)


class Benchmark(Generic[F]):
    def __init__(self, fn: F, line_arg: str, line_vals: List[Any], *, unit: str):
        assert unit in UNIT_SCALE and len(set(line_vals)) == len(line_vals) > 0
        self._fn = fn
        self._line_arg = line_arg
        self._line_vals = line_vals
        self._unit = unit
        self._configs: BENCH_CONFIG = []
        self._fn_params = inspect.signature(fn).parameters
        self._unit_scale = UNIT_SCALE[unit]
        assert line_arg in self._fn_params, (
            f"line_arg {line_arg!r} is not a parameter of {fn.__name__}; "
            f"available: {list(self._fn_params)}"
        )
        self._seen_args = {line_arg}

    def add_config(self, names: Tuple[str, ...], vals: List[Tuple[Any, ...]]) -> None:
        """Prepend a parametrize axis. Validates that names are real parameters
        of the benchmark fn, and rejects duplicates / collisions with line_arg."""
        assert len(names) > 0, "parametrize: must provide at least one name"
        for name in names:
            assert name in self._fn_params, (
                f"parametrize name {name!r} is not a parameter of "
                f"{self._fn.__name__}; available: {list(self._fn_params)}"
            )
            assert (
                name not in self._seen_args
            ), f"parametrize name {name!r} is already used"
            self._seen_args.add(name)
        self._configs.insert(0, (names, vals))

    def _collect_results(self) -> Tuple[List[List[float]], List[List[float]], bool]:
        axis_names = [n for n, _ in self._configs]
        axis_vals = [v for _, v in self._configs]
        results: List[List[float]] = []
        bandwidth_results: List[List[float]] = []
        should_log_bandwidth = False
        for system in self._line_vals:
            latencies: List[float] = []
            bandwidths: List[float] = []
            for combo in itertools.product(*axis_vals):
                kwargs: Dict[str, Any] = {self._line_arg: system}
                for names, values in zip(axis_names, combo):
                    kwargs.update(zip(names, values))
                try:
                    result = self._fn(**kwargs)
                except BenchSkip:
                    latencies.append(float("nan"))
                    if not DISABLE_LOG_BANDWIDTH:
                        bandwidths.append(float("nan"))
                    continue
                latencies.append(result.times[0] / self._unit_scale)
                if not DISABLE_LOG_BANDWIDTH and result.memory_footprint is not None:
                    should_log_bandwidth = True
                    bandwidths.append(
                        result.memory_footprint / (1024**3) / result.times[0]
                    )
            results.append(latencies)
            bandwidth_results.append(bandwidths)
        return results, bandwidth_results, should_log_bandwidth

    def run(self) -> None:
        # Pre-check: every required fn param must be covered.
        flat_names = [n for names, _ in self._configs for n in names]
        kinds = (
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
            inspect.Parameter.KEYWORD_ONLY,
        )
        missing = {
            n
            for n, p in self._fn_params.items()
            if p.default is inspect.Parameter.empty and p.kind in kinds
        } - (set(flat_names) | {self._line_arg})
        assert not missing, (
            f"parameters not parametrized for {self._fn.__name__}: "
            f"{sorted(missing)}"
        )

        results, bandwidths, should_log_bw = self._collect_results()

        table = Table()
        table.col(min_width=0, pad=0, align="<")  # id column (tight, left-aligned)
        for name in flat_names:
            table.col(name)
        table.sep()
        for system in self._line_vals:
            table.col(f"{system}({self._unit})", min_width=15)
        if should_log_bw:
            table.sep()
            for system in self._line_vals:
                table.col(f"{system}(GB/s)", min_width=15)

        axis_vals = [v for _, v in self._configs]
        for row_id, combo in enumerate(itertools.product(*axis_vals)):
            cells: List[Any] = [row_id]
            cells.extend(v for vt in combo for v in vt)
            cells.extend(table.format_latency(r[row_id]) for r in results)
            if should_log_bw:
                cells.extend(table.format_bandwidth(r[row_id]) for r in bandwidths)
            table.row(*cells)

        table.print()


def benchmark(line_arg: str, line_vals: List[Any], *, unit: str = "us"):
    def decorator(fn: F) -> Benchmark[F]:
        return Benchmark(fn, line_arg, line_vals, unit=unit)

    return decorator


def parametrize(names: str, vals: List[Any], ci_vals: Optional[List[Any]] = None):
    """Add a parametrize axis. Pytest-style:

    - Single name:   `parametrize("dim", [1024, 4096])`
    - Multiple names (correlated):
                     `parametrize("h,d", [(1, 64), (2, 128)])`

    For multi-name axes, each value must be a tuple/list of matching length.
    """
    name_tuple = tuple(n.strip() for n in names.split(","))
    assert all(name_tuple), f"parametrize: empty name in {names!r}"
    arity = len(name_tuple)

    def _normalize(vs: List[Any]) -> List[Tuple[Any, ...]]:
        if arity == 1:
            return [(v,) for v in vs]
        out: List[Tuple[Any, ...]] = []
        for v in vs:
            assert isinstance(
                v, (tuple, list)
            ), f"parametrize: multi-name values must be tuples, got {v!r}"
            t = tuple(v)
            assert (
                len(t) == arity
            ), f"parametrize: each value must have length {arity}, got {t!r}"
            out.append(t)
        return out

    def decorator(bench: Benchmark[F]) -> Benchmark[F]:
        chosen = ci_vals if (ci_vals is not None and is_in_ci()) else vals
        bench.add_config(name_tuple, _normalize(chosen))
        return bench

    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,
    *,
    input_args: Tuple[Any, ...] = (),
    input_kwargs: Dict[str, Any] = {},
    use_cuda_graph: bool = True,
    warmup_iters: int = 50,
    replay_iters: int = 1000,
    metrics: Tuple[Metric, ...] = (0.5, "avg"),
    stream: torch.cuda.Stream | None = None,
    # NOTE: should only clone the read args to avoid L2 cache effect in cuda graph
    graph_clone_args: Iterable[int] | Literal["all"] | None = "all",
    graph_clone_kwargs: Iterable[str] | Literal["all"] | None = "all",
    # NOTE: for memory-bandwidth profiling
    disable_log_bandwidth: bool = DISABLE_LOG_BANDWIDTH,
    memory_args: Iterable[Any] | Literal["all"] | None = "all",
    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
    :param input_kwargs: Keyword arguments to pass to the function
    :param use_cuda_graph: Whether to use CUDA graph for benchmarking
    :param warmup_iters: Number of warm-up iterations to run before benchmarking
    :param replay_iters: Number of iterations to run for benchmarking
    :param metrics: Metrics to compute from the timing results (quantiles in [0, 1] or "avg")
    :param stream: CUDA stream to use for benchmarking (if None, a new stream will be created)
    :param graph_clone_args: Indices of input_args to clone for each iteration.
                             Only the read args need to be cloned to avoid L2 cache effect.
    :param graph_clone_kwargs: Keys of input_kwargs to clone for each iteration.
                               Only the read args need to be cloned to avoid L2 cache effect.
    :param disable_log_bandwidth: Whether to disable logging memory bandwidth in the profile report.
    :param memory_args: Optional sequence of arguments to calculate total memory footprint.
                        Used for memory bandwidth estimation in the profile report.
    :param memory_output: Arguments whose output memory should be included in the memory footprint.
    :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()
    if stream is None:
        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
            if graph_clone_args == "all":
                graph_clone_args = range(len(input_args))
            elif graph_clone_args is None:
                graph_clone_args = []
            if graph_clone_kwargs == "all":
                graph_clone_kwargs = input_kwargs.keys()
            elif graph_clone_kwargs is None:
                graph_clone_kwargs = []
            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
            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)
                stream.synchronize()
                result.append(tic.elapsed_time(toc))

    stream.synchronize()
    result = _process_metrics(result, metrics)
    memory_footprint = None
    if not disable_log_bandwidth:
        if memory_args == "all":
            memory_args = input_args + tuple(input_kwargs.values())
        if memory_output == "out":
            memory_output = fn(*input_args, **input_kwargs)
        memory_footprint = extra_memory_footprint
        memory_footprint += _get_nbytes_recursive(extra_memory_args)
        memory_footprint += _get_nbytes_recursive(memory_args)
        memory_footprint += _get_nbytes_recursive(memory_output)

    return BenchResult(metrics, result, memory_footprint)