| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark sparse-indexer decode top-k on variable-length FP32 rows. |
| 4 | |
| 5 | Run on a reserved GPU, for example: |
| 6 | chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py |
| 7 | |
| 8 | Use --capacity to emulate graphs captured at a larger maximum context length. |
| 9 | Run the same command on the base and candidate builds to compare JSON results. |
| 10 | """ |
| 11 | |
| 12 | import argparse |
| 13 | import itertools |
| 14 | import json |
| 15 | import statistics |
| 16 | from pathlib import Path |
| 17 | |
| 18 | import torch |
| 19 | from flashinfer.testing import bench_gpu_time_with_cupti |
| 20 | |
| 21 | import vllm._custom_ops # noqa: F401 |
| 22 | from vllm.platforms import current_platform |
| 23 | |
| 24 | |
| 25 | def benchmark(rows: int, width: int, k: int, args) -> dict: |
| 26 | capacity = args.capacity or width |
| 27 | if capacity < width: |
| 28 | raise ValueError("Capacity must be at least the valid-length upper bound") |
| 29 | torch.manual_seed(args.seed) |
| 30 | logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32) |
| 31 | if args.full_length: |
| 32 | lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32) |
| 33 | else: |
| 34 | lengths = torch.randint( |
| 35 | int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32 |
| 36 | ) |
| 37 | invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None] |
| 38 | logits.masked_fill_(invalid, float("nan")) |
| 39 | out = torch.empty((rows, k), device="cuda", dtype=torch.int32) |
| 40 | workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8) |
| 41 | cooperative = ( |
| 42 | args.backend == "auto" |
| 43 | and rows <= 64 |
| 44 | and capacity % 4 == 0 |
| 45 | and current_platform.has_device_capability(90) |
| 46 | and not current_platform.is_device_capability_family(120) |
| 47 | ) |
| 48 | op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk |
| 49 | |
| 50 | def run(): |
| 51 | op(logits, lengths, out, workspace, k, capacity) |
| 52 | |
| 53 | run() |
| 54 | valid = torch.arange(k, device="cuda")[None] < lengths[:, None] |
| 55 | assert torch.all(out[~valid] == -1) |
| 56 | assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid) |
| 57 | sorted_indices = out.sort(dim=1).values |
| 58 | assert torch.all( |
| 59 | (sorted_indices[:, 1:] != sorted_indices[:, :-1]) |
| 60 | | (sorted_indices[:, 1:] == -1) |
| 61 | ) |
| 62 | selected = logits.gather(1, out.clamp_min(0).long()).masked_fill( |
| 63 | ~valid, -float("inf") |
| 64 | ) |
| 65 | reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values |
| 66 | torch.testing.assert_close( |
| 67 | selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0 |
| 68 | ) |
| 69 | trials = [ |
| 70 | statistics.median( |
| 71 | bench_gpu_time_with_cupti( |
| 72 | run, |
| 73 | use_cuda_graph=True, |
| 74 | cold_l2_cache=not args.warm, |
| 75 | dry_run_iters=5, |
| 76 | repeat_iters=args.repeat, |
| 77 | ) |
| 78 | ) |
| 79 | * 1000 |
| 80 | for _ in range(args.trials) |
| 81 | ] |
| 82 | us = statistics.median(trials) |
| 83 | return dict( |
| 84 | rows=rows, |
| 85 | width=width, |
| 86 | capacity=capacity, |
| 87 | k=k, |
| 88 | backend="cooperative" if cooperative else "persistent", |
| 89 | us=us, |
| 90 | trials_us=trials, |
| 91 | # Minimum useful traffic: one valid-score read and one index write. |
| 92 | effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000), |
| 93 | ) |
| 94 | |
| 95 | |
| 96 | def main(): |
| 97 | parser = argparse.ArgumentParser(description=__doc__) |
| 98 | parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024]) |
| 99 | parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840]) |
| 100 | parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048]) |
| 101 | parser.add_argument("--capacity", type=int) |
| 102 | parser.add_argument( |
| 103 | "--full-length", action="store_true", help="Use exactly --widths valid scores" |
| 104 | ) |
| 105 | parser.add_argument("--backend", choices=["auto", "persistent"], default="auto") |
| 106 | parser.add_argument( |
| 107 | "--warm", action="store_true", help="Keep L2 warm between replays" |
| 108 | ) |
| 109 | parser.add_argument("--repeat", type=int, default=60) |
| 110 | parser.add_argument("--trials", type=int, default=3) |
| 111 | parser.add_argument("--seed", type=int, default=123) |
| 112 | parser.add_argument("--output", type=Path) |
| 113 | args = parser.parse_args() |
| 114 | metadata = dict( |
| 115 | gpu=str(torch.cuda.get_device_properties(0)), |
| 116 | torch=torch.__version__, |
| 117 | cuda=torch.version.cuda, |
| 118 | args=vars(args), |
| 119 | ) |
| 120 | print(json.dumps(metadata, default=str)) |
| 121 | results = [] |
| 122 | for rows, width, k in itertools.product(args.rows, args.widths, args.top_k): |
| 123 | result = benchmark(rows, width, k, args) |
| 124 | results.append(result) |
| 125 | print(json.dumps(result), flush=True) |
| 126 | if args.output: |
| 127 | args.output.write_text( |
| 128 | json.dumps(dict(metadata=metadata, results=results), indent=2, default=str) |
| 129 | + "\n" |
| 130 | ) |
| 131 | |
| 132 | |
| 133 | if __name__ == "__main__": |
| 134 | main() |