@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 158 lines 4.72 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark the ROCm sparse-indexer decode top-k in isolation.
4
5 Example usage:
6 python3 benchmarks/kernels/benchmark_top_k_per_row_decode.py \
7 --rows 4 16 32 --lens 2048 8192 16384 65536 --capacity 262144
8 """
9
10 import argparse
11 import itertools
12 import json
13
14 import torch
15
16 import vllm._custom_ops # noqa: F401
17 from vllm._aiter_ops import rocm_aiter_ops
18 from vllm.triton_utils import triton
19
20 DISTRIBUTIONS = ("indexer", "fp8", "randn")
21
22
23 def make_logits(rows: int, capacity: int, dist: str, spread: float):
24 normal = torch.randn(rows, capacity, device="cuda", dtype=torch.float32)
25 if dist == "randn":
26 return normal
27 if dist == "fp8":
28 return normal.to(torch.float8_e4m3fnuz).to(torch.float32)
29 return normal * spread + 3.0
30
31
32 def make_inputs(rows: int, live: int, capacity: int, k: int, args):
33 torch.manual_seed(args.seed)
34 logits = make_logits(rows, capacity, args.dist, args.spread)
35 seq_lens = torch.full((rows,), live, device="cuda", dtype=torch.int32)
36 invalid = torch.arange(capacity, device="cuda")[None] >= seq_lens[:, None]
37 logits.masked_fill_(invalid, -float("inf"))
38 indices = torch.empty((rows, k), device="cuda", dtype=torch.int32)
39 return logits, seq_lens, indices
40
41
42 def run_native(logits, seq_lens, indices, k):
43 torch.ops._C.top_k_per_row_decode(
44 logits,
45 1,
46 seq_lens,
47 indices,
48 logits.shape[0],
49 logits.stride(0),
50 logits.stride(1),
51 k,
52 )
53
54
55 def run_aiter(logits, seq_lens, indices, k):
56 rocm_aiter_ops.indexer_top_k_decode(logits, 1, seq_lens, indices, k)
57
58
59 BACKENDS = {"native": run_native, "aiter": run_aiter}
60
61
62 def check(logits, indices, k):
63 selected = logits.gather(1, indices.long()).sort(dim=1, descending=True).values
64 expected = logits.topk(k, dim=1).values
65 torch.testing.assert_close(selected, expected, atol=0, rtol=0)
66
67
68 def capture(run):
69 side = torch.cuda.Stream()
70 side.wait_stream(torch.cuda.current_stream())
71 with torch.cuda.stream(side):
72 for _ in range(3):
73 run()
74 torch.cuda.current_stream().wait_stream(side)
75 graph = torch.cuda.CUDAGraph()
76 with torch.cuda.graph(graph):
77 run()
78 return graph
79
80
81 def benchmark(backend: str, rows: int, live: int, args) -> dict:
82 logits, seq_lens, indices = make_inputs(rows, live, args.capacity, args.top_k, args)
83 fn = BACKENDS[backend]
84
85 def run():
86 fn(logits, seq_lens, indices, args.top_k)
87
88 run()
89 torch.accelerator.synchronize()
90 check(logits, indices, args.top_k)
91 timed = run
92 if args.graph:
93 graph = capture(run)
94 indices.zero_()
95 graph.replay()
96 torch.accelerator.synchronize()
97 check(logits, indices, args.top_k)
98 timed = graph.replay
99 ms = triton.testing.do_bench(timed, warmup=args.warmup, rep=args.rep)
100 return dict(
101 backend=backend,
102 mode="graph" if args.graph else "eager",
103 dist=args.dist,
104 rows=rows,
105 live=live,
106 capacity=args.capacity,
107 k=args.top_k,
108 us=ms * 1000,
109 live_gbps=rows * live * 4 / (ms * 1e6),
110 )
111
112
113 def main():
114 parser = argparse.ArgumentParser(description=__doc__)
115 parser.add_argument("--rows", type=int, nargs="+", default=[1, 2, 4, 8])
116 parser.add_argument(
117 "--lens",
118 type=int,
119 nargs="+",
120 default=[32768, 65536, 98304, 131072],
121 help="Valid columns per row; must be at least --top-k",
122 )
123 parser.add_argument("--capacity", type=int, default=262144)
124 parser.add_argument("--top-k", type=int, default=512)
125 parser.add_argument(
126 "--backends", nargs="+", choices=list(BACKENDS), default=["native", "aiter"]
127 )
128 parser.add_argument(
129 "--graph",
130 action="store_true",
131 help="Time a captured graph replay, matching how decode actually runs",
132 )
133 parser.add_argument(
134 "--dist",
135 choices=DISTRIBUTIONS,
136 default="indexer",
137 help="Logit value distribution; see the module docstring",
138 )
139 parser.add_argument(
140 "--spread",
141 type=float,
142 default=0.01,
143 help="Relative width of the --dist indexer cluster",
144 )
145 parser.add_argument("--warmup", type=int, default=25)
146 parser.add_argument("--rep", type=int, default=200)
147 parser.add_argument("--seed", type=int, default=123)
148 args = parser.parse_args()
149
150 print(json.dumps(dict(gpu=torch.cuda.get_device_name(0), args=vars(args))))
151 for backend, rows, live in itertools.product(args.backends, args.rows, args.lens):
152 if live < args.top_k:
153 continue
154 print(json.dumps(benchmark(backend, rows, live, args)), flush=True)
155
156
157 if __name__ == "__main__":
158 main()