chore: import upstream snapshot with attribution

python/sglang/jit_kernel/benchmark/kv_canary/utils.py

@@ -0,0 +1,404 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource
+
+BS_AXIS: list[int] = [1, 4, 32, 128, 256, 1024]
+PREFIX_AXIS: list[int] = [0, 128, 1024, 4096, 10240, 16384]
+EXTEND_LEN_AXIS: list[int] = [128, 512, 4096, 16384]
+POOL_AXIS: list[str] = ["full", "swa_window_128"]
+REAL_KV_AXIS: list[str] = ["none", "small_1src", "med_2src", "max_4src"]
+HASH_MODE_AXIS: list[str] = ["none", "partial", "all"]
+SWA_WINDOW: int = 128
+RING_CAPACITY: int = 256
+MAX_EXTEND_TOKENS_PER_FORWARD: int = 4096
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class BenchCase:
+    scenario: str
+    bs: int
+    prefix_len: int
+    mode: str
+    extend_len: int
+    pool_kind: str
+    real_kv_kind: str
+    hash_mode: str
+
+    @property
+    def case_id(self) -> str:
+        return (
+            f"{self.scenario}_bs{self.bs}_prefix{self.prefix_len}_{self.mode}{self.extend_len}"
+            f"_{self.pool_kind}_rkv{self.real_kv_kind}_hash{self.hash_mode}"
+        )
+
+
+def _case(
+    *,
+    scenario: str,
+    bs: int,
+    prefix_len: int,
+    mode: str,
+    extend_len: int,
+    pool_kind: str,
+    real_kv_kind: str = "none",
+    hash_mode: str = "none",
+) -> BenchCase:
+    return BenchCase(
+        scenario=scenario,
+        bs=bs,
+        prefix_len=prefix_len,
+        mode=mode,
+        extend_len=extend_len,
+        pool_kind=pool_kind,
+        real_kv_kind=real_kv_kind,
+        hash_mode=hash_mode,
+    )
+
+
+def _is_realistic_extend_case(case: BenchCase) -> bool:
+    if case.mode != "extend":
+        return True
+    return case.bs * case.extend_len <= MAX_EXTEND_TOKENS_PER_FORWARD
+
+
+def _dedupe_cases(cases: list[BenchCase]) -> list[BenchCase]:
+    seen: set[str] = set()
+    result: list[BenchCase] = []
+
+    for case in cases:
+        if case.case_id in seen:
+            continue
+        seen.add(case.case_id)
+        result.append(case)
+
+    return result
+
+
+def build_fast_matrix_cases() -> list[BenchCase]:
+    return _dedupe_cases(
+        [
+            _case(
+                scenario="smoke_decode_empty",
+                bs=1,
+                prefix_len=0,
+                mode="decode",
+                extend_len=1,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="small_extend_batch",
+                bs=32,
+                prefix_len=4096,
+                mode="extend",
+                extend_len=128,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="e2e_decode_steady",
+                bs=256,
+                prefix_len=4096,
+                mode="decode",
+                extend_len=1,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="decode_large_batch_short_prefix",
+                bs=1024,
+                prefix_len=1024,
+                mode="decode",
+                extend_len=1,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="e2e_prefill_chunk_first",
+                bs=1,
+                prefix_len=0,
+                mode="extend",
+                extend_len=4096,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="e2e_prefill_chunk_mid",
+                bs=1,
+                prefix_len=8192,
+                mode="extend",
+                extend_len=4096,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="e2e_prefill_chunk_last",
+                bs=1,
+                prefix_len=12288,
+                mode="extend",
+                extend_len=4096,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="e2e_decode_tail",
+                bs=1,
+                prefix_len=5120,
+                mode="decode",
+                extend_len=1,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="swa_decode_long_prefix",
+                bs=128,
+                prefix_len=10240,
+                mode="decode",
+                extend_len=1,
+                pool_kind="swa_window_128",
+            ),
+            _case(
+                scenario="small_extend_single_req",
+                bs=1,
+                prefix_len=128,
+                mode="extend",
+                extend_len=128,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="medium_extend_chunk",
+                bs=4,
+                prefix_len=1024,
+                mode="extend",
+                extend_len=512,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="decode_mid_batch",
+                bs=128,
+                prefix_len=4096,
+                mode="decode",
+                extend_len=1,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="e2e_prefill_chunk_second",
+                bs=1,
+                prefix_len=4096,
+                mode="extend",
+                extend_len=4096,
+                pool_kind="full",
+            ),
+            _case(
+                scenario="swa_decode_short_prefix",
+                bs=256,
+                prefix_len=128,
+                mode="decode",
+                extend_len=1,
+                pool_kind="swa_window_128",
+            ),
+            _case(
+                scenario="swa_decode_tail",
+                bs=4,
+                prefix_len=10240,
+                mode="decode",
+                extend_len=1,
+                pool_kind="swa_window_128",
+            ),
+            _case(
+                scenario="small_extend_batch_hash",
+                bs=32,
+                prefix_len=4096,
+                mode="extend",
+                extend_len=128,
+                pool_kind="full",
+                real_kv_kind="small_1src",
+                hash_mode="partial",
+            ),
+            _case(
+                scenario="e2e_prefill_chunk_hash",
+                bs=1,
+                prefix_len=12288,
+                mode="extend",
+                extend_len=4096,
+                pool_kind="full",
+                real_kv_kind="med_2src",
+                hash_mode="all",
+            ),
+            _case(
+                scenario="e2e_decode_steady_hash",
+                bs=256,
+                prefix_len=4096,
+                mode="decode",
+                extend_len=1,
+                pool_kind="full",
+                real_kv_kind="max_4src",
+                hash_mode="all",
+            ),
+            _case(
+                scenario="swa_decode_long_prefix_hash",
+                bs=128,
+                prefix_len=10240,
+                mode="decode",
+                extend_len=1,
+                pool_kind="swa_window_128",
+                real_kv_kind="med_2src",
+                hash_mode="partial",
+            ),
+            _case(
+                scenario="smoke_decode_empty_hash",
+                bs=1,
+                prefix_len=0,
+                mode="decode",
+                extend_len=1,
+                pool_kind="full",
+                real_kv_kind="small_1src",
+                hash_mode="all",
+            ),
+        ]
+    )
+
+
+def build_full_matrix_cases() -> list[BenchCase]:
+    """Full matrix plus targeted e2e points.
+
+    Extend cases are pruned to a maximum token chunk per forward because the scheduler chunks long
+    prefills; for example, a 4096-token extend is represented as ``bs=1``, not ``bs=32``.
+    """
+    fast = build_fast_matrix_cases()
+    fast_keys = {c.case_id for c in fast}
+    full: list[BenchCase] = list(fast)
+
+    for bs in BS_AXIS:
+        for prefix_len in PREFIX_AXIS:
+            for pool_kind in POOL_AXIS:
+                for mode, extend_len in (
+                    ("decode", 1),
+                    *(("extend", e) for e in EXTEND_LEN_AXIS),
+                ):
+                    case = _case(
+                        scenario="matrix",
+                        bs=bs,
+                        prefix_len=prefix_len,
+                        mode=mode,
+                        extend_len=extend_len,
+                        pool_kind=pool_kind,
+                    )
+                    if not _is_realistic_extend_case(case):
+                        continue
+                    if case.case_id in fast_keys:
+                        continue
+                    full.append(case)
+
+    fast_base_points = [
+        (c.bs, c.prefix_len, c.mode, c.extend_len, c.pool_kind)
+        for c in fast
+        if c.real_kv_kind == "none" and c.hash_mode == "none"
+    ]
+    for bs, prefix_len, mode, extend_len, pool_kind in fast_base_points:
+        for hash_mode in HASH_MODE_AXIS:
+            if hash_mode == "none":
+                continue
+            for real_kv_kind in REAL_KV_AXIS:
+                if real_kv_kind == "none":
+                    continue
+                case = _case(
+                    scenario="fold_matrix",
+                    bs=bs,
+                    prefix_len=prefix_len,
+                    mode=mode,
+                    extend_len=extend_len,
+                    pool_kind=pool_kind,
+                    real_kv_kind=real_kv_kind,
+                    hash_mode=hash_mode,
+                )
+                if not _is_realistic_extend_case(case):
+                    continue
+                if case.case_id in fast_keys:
+                    continue
+                full.append(case)
+                fast_keys.add(case.case_id)
+
+    return full
+
+
+def cases_to_x_vals(
+    cases: list[BenchCase],
+) -> list[tuple[str, int, int, str, int, str, str, str]]:
+    return [
+        (
+            c.scenario,
+            c.bs,
+            c.prefix_len,
+            c.mode,
+            c.extend_len,
+            c.pool_kind,
+            c.real_kv_kind,
+            c.hash_mode,
+        )
+        for c in cases
+    ]
+
+
+def _one_real_kv_source(
+    *, num_slots: int, num_bytes: int, read_bytes: int, device: torch.device
+) -> RealKvSource:
+    tensor = torch.zeros(max(1, num_slots), num_bytes, dtype=torch.uint8, device=device)
+    return RealKvSource(
+        tensor=tensor,
+        page_size=1,
+        num_bytes_per_token=num_bytes,
+        read_bytes=read_bytes,
+    )
+
+
+def make_real_kv_sources(
+    *, kind: str, num_slots: int, device: torch.device
+) -> tuple[RealKvSource, ...]:
+    """Map a ``real_kv_kind`` axis label to a tuple of ``RealKvSource`` configs.
+
+    Byte-volume ladder (none -> small_1src -> med_2src -> max_4src) so the bench exposes the
+    ``real_kv_fold_sources`` PARTIAL/ALL cost gradient. ``max_4src`` hits the
+    ``consts.MAX_REAL_KV_SOURCES = 4`` ABI ceiling.
+    """
+    if kind == "none":
+        return ()
+    if kind == "small_1src":
+        return (
+            _one_real_kv_source(
+                num_slots=num_slots, num_bytes=16, read_bytes=16, device=device
+            ),
+        )
+    if kind == "med_2src":
+        return tuple(
+            _one_real_kv_source(
+                num_slots=num_slots, num_bytes=32, read_bytes=16, device=device
+            )
+            for _ in range(2)
+        )
+    if kind == "max_4src":
+        return tuple(
+            _one_real_kv_source(
+                num_slots=num_slots, num_bytes=64, read_bytes=32, device=device
+            )
+            for _ in range(4)
+        )
+    raise ValueError(f"kv-canary bench: unknown real_kv_kind {kind!r}")
+
+
+def naive_slot_copy_fn(*, total: int, device: torch.device) -> Callable[[], None]:
+    n_slots = max(total, 1)
+    payload = torch.zeros(n_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device)
+    sink = torch.zeros_like(payload)
+    indices = torch.arange(n_slots, device=device, dtype=torch.int64) % sink.shape[0]
+
+    def baseline() -> None:
+        sink.index_copy_(0, indices, payload)
+
+    return baseline
+
+
+def naive_cumsum_fn(*, bs: int, device: torch.device) -> Callable[[], None]:
+    counts = torch.zeros(max(bs, 1), dtype=torch.int32, device=device)
+
+    def baseline() -> None:
+        torch.cumsum(counts, dim=0)
+
+    return baseline

python/sglang/jit_kernel/benchmark/marker.py

@@ -0,0 +1,554 @@
+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)

python/sglang/jit_kernel/benchmark/utils.py

@@ -0,0 +1,108 @@
+"""Common utilities for jit_kernel benchmark files."""
+
+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"
+DEFAULT_QUANTILES = [0.5, 0.2, 0.8]
+
+
+def create_empty(*shape: int, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE):
+    return torch.empty(shape, dtype=dtype, device=device)
+
+
+def create_random(*shape: int, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE):
+    return torch.randn(shape, dtype=dtype, device=device)
+
+
+def get_benchmark_range(full_range: List, ci_range: List) -> List:
+    """Return appropriate benchmark range based on CI environment."""
+    return ci_range if is_in_ci() else full_range
+
+
+def run_benchmark(
+    fn: Callable,
+    quantiles: Sequence[float] = (),
+    scale: float = 1.0,
+) -> Tuple[float, float, float]:
+    """Execute benchmark using CUDA graph and return times in microseconds.
+
+    Args:
+        fn: Function to benchmark
+        quantiles: Quantiles for timing measurements [median, min, max]
+        scale: Scale the result down (usually num_layers).
+
+    Returns:
+        Tuple of (median_us, max_us, min_us)
+    """
+    quantiles = list(quantiles or DEFAULT_QUANTILES)
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return 1000 * ms / scale, 1000 * max_ms / scale, 1000 * min_ms / scale
+
+
+def run_benchmark_no_cudagraph(
+    fn: Callable,
+    quantiles: Sequence[float] = (),
+    scale: float = 1.0,
+) -> Tuple[float, float, float]:
+    quantiles = list(quantiles or DEFAULT_QUANTILES)
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+    return 1000 * ms / scale, 1000 * max_ms / scale, 1000 * min_ms / scale