| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import statistics |
| 5 | |
| 6 | import torch |
| 7 | from tabulate import tabulate |
| 8 | |
| 9 | from vllm.models.inkling.nvidia.ops import qkvr_prep |
| 10 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 11 | |
| 12 | |
| 13 | def make_inputs(tokens: int, tp_size: int, is_local: bool): |
| 14 | torch.manual_seed(0) |
| 15 | num_q_heads = 64 // tp_size |
| 16 | num_kv_heads = (16 if is_local else 8) // tp_size |
| 17 | head_dim = 128 |
| 18 | d_rel = 16 |
| 19 | rel_extent = 512 if is_local else 1024 |
| 20 | page_size = 16 |
| 21 | num_blocks = (tokens + page_size - 1) // page_size |
| 22 | q_width = num_q_heads * head_dim |
| 23 | kv_width = num_kv_heads * head_dim |
| 24 | r_width = num_q_heads * d_rel |
| 25 | device = "cuda" |
| 26 | |
| 27 | qkvr = torch.randn( |
| 28 | tokens, |
| 29 | q_width + 2 * kv_width + r_width, |
| 30 | device=device, |
| 31 | dtype=torch.bfloat16, |
| 32 | ) |
| 33 | k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16) |
| 34 | v_weight = torch.randn_like(k_weight) |
| 35 | q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16) |
| 36 | k_norm_weight = torch.randn_like(q_norm_weight) |
| 37 | rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16) |
| 38 | conv_cache = torch.zeros( |
| 39 | num_blocks, |
| 40 | num_kv_heads, |
| 41 | page_size, |
| 42 | 2 * head_dim, |
| 43 | device=device, |
| 44 | dtype=torch.bfloat16, |
| 45 | ) |
| 46 | key_cache = torch.empty( |
| 47 | num_blocks, |
| 48 | page_size, |
| 49 | num_kv_heads, |
| 50 | head_dim, |
| 51 | device=device, |
| 52 | dtype=torch.bfloat16, |
| 53 | ) |
| 54 | value_cache = torch.empty_like(key_cache) |
| 55 | positions = torch.arange(tokens, device=device, dtype=torch.int64) |
| 56 | block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None] |
| 57 | seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32) |
| 58 | slots = torch.arange(tokens, device=device, dtype=torch.int64) |
| 59 | query_start = torch.zeros(tokens, device=device, dtype=torch.int32) |
| 60 | log_scaling = None |
| 61 | if not is_local: |
| 62 | effective_n = (positions + 1).to(torch.float32) |
| 63 | log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0)) |
| 64 | return ( |
| 65 | qkvr, |
| 66 | k_weight, |
| 67 | v_weight, |
| 68 | q_norm_weight, |
| 69 | k_norm_weight, |
| 70 | rel_proj, |
| 71 | 1e-6, |
| 72 | num_q_heads, |
| 73 | num_kv_heads, |
| 74 | head_dim, |
| 75 | d_rel, |
| 76 | conv_cache, |
| 77 | key_cache, |
| 78 | value_cache, |
| 79 | positions, |
| 80 | block_table, |
| 81 | seq_idx, |
| 82 | slots, |
| 83 | query_start, |
| 84 | slots, |
| 85 | 0, |
| 86 | head_dim, |
| 87 | page_size, |
| 88 | log_scaling, |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def capture(implementation, inputs): |
| 93 | outputs = [] |
| 94 | |
| 95 | def run(): |
| 96 | outputs[:] = implementation.fused_qkvr_prep(*inputs) |
| 97 | |
| 98 | stream = torch.cuda.Stream() |
| 99 | stream.wait_stream(torch.cuda.current_stream()) |
| 100 | with torch.cuda.stream(stream): |
| 101 | for _ in range(3): |
| 102 | run() |
| 103 | torch.cuda.current_stream().wait_stream(stream) |
| 104 | torch.accelerator.synchronize() |
| 105 | graph = torch.cuda.CUDAGraph() |
| 106 | with torch.cuda.graph(graph): |
| 107 | run() |
| 108 | torch.accelerator.synchronize() |
| 109 | return graph, outputs |
| 110 | |
| 111 | |
| 112 | def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float: |
| 113 | for _ in range(warmup): |
| 114 | graph.replay() |
| 115 | torch.accelerator.synchronize() |
| 116 | start = torch.cuda.Event(enable_timing=True) |
| 117 | end = torch.cuda.Event(enable_timing=True) |
| 118 | start.record() |
| 119 | for _ in range(repeats): |
| 120 | graph.replay() |
| 121 | end.record() |
| 122 | end.synchronize() |
| 123 | return start.elapsed_time(end) * 1000 / repeats |
| 124 | |
| 125 | |
| 126 | def benchmark(inputs, args) -> float: |
| 127 | graph, _ = capture(qkvr_prep, inputs) |
| 128 | return statistics.median( |
| 129 | time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials) |
| 130 | ) |
| 131 | |
| 132 | |
| 133 | @torch.inference_mode() |
| 134 | def main(args): |
| 135 | rows = [] |
| 136 | for tp_size in args.tp_sizes: |
| 137 | for tokens in args.tokens: |
| 138 | for is_local in (True, False): |
| 139 | triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args) |
| 140 | rows.append( |
| 141 | [ |
| 142 | tp_size, |
| 143 | tokens, |
| 144 | "local" if is_local else "global", |
| 145 | triton_us, |
| 146 | ] |
| 147 | ) |
| 148 | |
| 149 | print("Inkling QKVR prep (CUDA graph, median latency)") |
| 150 | print( |
| 151 | tabulate( |
| 152 | rows, |
| 153 | headers=[ |
| 154 | "TP", |
| 155 | "tokens", |
| 156 | "scope", |
| 157 | "Triton (us)", |
| 158 | ], |
| 159 | floatfmt=("d", "d", "", ".2f"), |
| 160 | ) |
| 161 | ) |
| 162 | |
| 163 | |
| 164 | if __name__ == "__main__": |
| 165 | parser = FlexibleArgumentParser() |
| 166 | parser.add_argument( |
| 167 | "--tokens", |
| 168 | type=int, |
| 169 | nargs="+", |
| 170 | default=[1 << power for power in range(15)], |
| 171 | ) |
| 172 | parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8]) |
| 173 | parser.add_argument("--warmup", type=int, default=20) |
| 174 | parser.add_argument("--repeats", type=int, default=200) |
| 175 | parser.add_argument("--trials", type=int, default=5) |
| 176 | main(parser.parse_args()) |