| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from collections.abc import Callable |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm import _custom_ops as ops |
| 10 | from vllm.triton_utils import triton |
| 11 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 12 | |
| 13 | SCENARIOS = { |
| 14 | "single-60k": [60_000], |
| 15 | "single-300k": [300_000], |
| 16 | "skew-2": [60_000, 300_000], |
| 17 | "skew-4": [60_000, 100_000, 180_000, 300_000], |
| 18 | "skew-8": [ |
| 19 | 60_000, |
| 20 | 60_000, |
| 21 | 80_000, |
| 22 | 100_000, |
| 23 | 140_000, |
| 24 | 180_000, |
| 25 | 240_000, |
| 26 | 300_000, |
| 27 | ], |
| 28 | } |
| 29 | DTYPES = { |
| 30 | "fp8": torch.float8_e4m3fn, |
| 31 | "bfloat16": torch.bfloat16, |
| 32 | "float32": torch.float32, |
| 33 | } |
| 34 | |
| 35 | |
| 36 | def make_page_table( |
| 37 | seq_lens: list[int], |
| 38 | block_size: int, |
| 39 | seq_starts: list[int] | None = None, |
| 40 | ) -> tuple[torch.Tensor, torch.Tensor, int]: |
| 41 | if seq_starts is None: |
| 42 | seq_starts = [0] * len(seq_lens) |
| 43 | blocks_per_req = [ |
| 44 | math.ceil((start + length) / block_size) |
| 45 | for start, length in zip(seq_starts, seq_lens) |
| 46 | ] |
| 47 | total_blocks = sum(blocks_per_req) |
| 48 | block_table = torch.zeros( |
| 49 | (len(seq_lens), max(blocks_per_req)), dtype=torch.int32, device="cuda" |
| 50 | ) |
| 51 | physical_blocks = torch.randperm(total_blocks, dtype=torch.int32, device="cuda") |
| 52 | block_offset = 0 |
| 53 | for req_id, num_blocks in enumerate(blocks_per_req): |
| 54 | block_table[req_id, :num_blocks] = physical_blocks[ |
| 55 | block_offset : block_offset + num_blocks |
| 56 | ] |
| 57 | block_offset += num_blocks |
| 58 | |
| 59 | cu_seq_lens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device="cuda") |
| 60 | cu_seq_lens[1:] = torch.tensor(seq_lens, dtype=torch.int32, device="cuda").cumsum( |
| 61 | dim=0 |
| 62 | ) |
| 63 | return block_table, cu_seq_lens, total_blocks |
| 64 | |
| 65 | |
| 66 | def make_cache_gather( |
| 67 | seq_lens: list[int], |
| 68 | block_size: int, |
| 69 | entry_size: int, |
| 70 | dtype: torch.dtype, |
| 71 | ) -> tuple[Callable[[], None], int]: |
| 72 | seq_starts_list = [ |
| 73 | 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens)) |
| 74 | ] |
| 75 | block_table, cu_seq_lens, total_blocks = make_page_table( |
| 76 | seq_lens, block_size, seq_starts_list |
| 77 | ) |
| 78 | src_cache = torch.empty( |
| 79 | (total_blocks, block_size, entry_size), dtype=dtype, device="cuda" |
| 80 | ) |
| 81 | dst = torch.empty((sum(seq_lens), entry_size), dtype=dtype, device="cuda") |
| 82 | seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda") |
| 83 | |
| 84 | def run() -> None: |
| 85 | ops.cp_gather_cache( |
| 86 | src_cache, |
| 87 | dst, |
| 88 | block_table, |
| 89 | cu_seq_lens, |
| 90 | len(seq_lens), |
| 91 | seq_starts, |
| 92 | ) |
| 93 | |
| 94 | bytes_moved = 2 * dst.numel() * dst.element_size() |
| 95 | return run, bytes_moved |
| 96 | |
| 97 | |
| 98 | def make_fp8_upconvert( |
| 99 | seq_lens: list[int], |
| 100 | block_size: int, |
| 101 | ) -> tuple[Callable[[], None], int]: |
| 102 | entry_bytes = 656 |
| 103 | output_elements = 576 |
| 104 | seq_starts_list = [ |
| 105 | 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens)) |
| 106 | ] |
| 107 | block_table, cu_seq_lens, total_blocks = make_page_table( |
| 108 | seq_lens, block_size, seq_starts_list |
| 109 | ) |
| 110 | src_cache = torch.empty( |
| 111 | (total_blocks, block_size, entry_bytes), |
| 112 | dtype=torch.uint8, |
| 113 | device="cuda", |
| 114 | ) |
| 115 | dst = torch.empty( |
| 116 | (sum(seq_lens), output_elements), dtype=torch.bfloat16, device="cuda" |
| 117 | ) |
| 118 | seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda") |
| 119 | |
| 120 | def run() -> None: |
| 121 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 122 | src_cache, |
| 123 | dst, |
| 124 | block_table, |
| 125 | cu_seq_lens[:-1], |
| 126 | len(seq_lens), |
| 127 | seq_starts, |
| 128 | ) |
| 129 | |
| 130 | bytes_moved = sum(seq_lens) * (entry_bytes + output_elements * 2) |
| 131 | return run, bytes_moved |
| 132 | |
| 133 | |
| 134 | def make_maybe_dequant_gather( |
| 135 | seq_lens: list[int], |
| 136 | block_size: int, |
| 137 | entry_size: int, |
| 138 | ) -> tuple[Callable[[], None], int]: |
| 139 | seq_starts_list = [ |
| 140 | 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens)) |
| 141 | ] |
| 142 | block_table, cu_seq_lens, total_blocks = make_page_table( |
| 143 | seq_lens, block_size, seq_starts_list |
| 144 | ) |
| 145 | src_cache = torch.empty( |
| 146 | (total_blocks, block_size, entry_size), |
| 147 | dtype=torch.float8_e4m3fn, |
| 148 | device="cuda", |
| 149 | ) |
| 150 | dst = torch.empty((sum(seq_lens), entry_size), dtype=torch.bfloat16, device="cuda") |
| 151 | token_to_seq = torch.repeat_interleave( |
| 152 | torch.arange(len(seq_lens), dtype=torch.int32, device="cuda"), |
| 153 | torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), |
| 154 | ) |
| 155 | seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda") |
| 156 | scale = torch.tensor(0.1, dtype=torch.float32, device="cuda") |
| 157 | |
| 158 | def run() -> None: |
| 159 | ops.gather_and_maybe_dequant_cache( |
| 160 | src_cache, |
| 161 | dst, |
| 162 | block_table, |
| 163 | cu_seq_lens, |
| 164 | token_to_seq, |
| 165 | sum(seq_lens), |
| 166 | "fp8", |
| 167 | scale, |
| 168 | seq_starts, |
| 169 | ) |
| 170 | |
| 171 | bytes_moved = sum(seq_lens) * entry_size * 3 |
| 172 | return run, bytes_moved |
| 173 | |
| 174 | |
| 175 | @torch.inference_mode() |
| 176 | def run_scenario( |
| 177 | variant: str, |
| 178 | name: str, |
| 179 | seq_lens: list[int], |
| 180 | block_size: int, |
| 181 | entry_size: int, |
| 182 | dtype: torch.dtype, |
| 183 | warmup_ms: int, |
| 184 | rep_ms: int, |
| 185 | ) -> None: |
| 186 | if variant == "cache": |
| 187 | run, bytes_moved = make_cache_gather(seq_lens, block_size, entry_size, dtype) |
| 188 | elif variant == "fp8-upconvert": |
| 189 | run, bytes_moved = make_fp8_upconvert(seq_lens, block_size) |
| 190 | else: |
| 191 | run, bytes_moved = make_maybe_dequant_gather(seq_lens, block_size, entry_size) |
| 192 | |
| 193 | latency_ms = triton.testing.do_bench( |
| 194 | run, warmup=warmup_ms, rep=rep_ms, return_mode="median" |
| 195 | ) |
| 196 | bandwidth_gbps = bytes_moved / latency_ms / 1e6 |
| 197 | lengths = ",".join(str(seq_len) for seq_len in seq_lens) |
| 198 | print( |
| 199 | f"{variant:15s} {name:10s} batch={len(seq_lens):2d} " |
| 200 | f"total={sum(seq_lens):7d} latency={latency_ms * 1e3:9.2f} us " |
| 201 | f"bandwidth={bandwidth_gbps:8.1f} GB/s lengths=[{lengths}]" |
| 202 | ) |
| 203 | |
| 204 | |
| 205 | def main() -> None: |
| 206 | parser = FlexibleArgumentParser(description="Benchmark cp_gather variants") |
| 207 | parser.add_argument( |
| 208 | "--variant", |
| 209 | choices=["all", "cache", "fp8-upconvert", "maybe-dequant"], |
| 210 | default="all", |
| 211 | ) |
| 212 | parser.add_argument("--scenario", choices=["all", *SCENARIOS], default="all") |
| 213 | parser.add_argument("--dtype", choices=DTYPES, default="fp8") |
| 214 | parser.add_argument("--block-size", type=int, default=64) |
| 215 | parser.add_argument("--entry-size", type=int, default=576) |
| 216 | parser.add_argument("--warmup-ms", type=int, default=25) |
| 217 | parser.add_argument("--rep-ms", type=int, default=100) |
| 218 | parser.add_argument("--seed", type=int, default=0) |
| 219 | args = parser.parse_args() |
| 220 | |
| 221 | torch.manual_seed(args.seed) |
| 222 | variants = ( |
| 223 | ("cache", "fp8-upconvert", "maybe-dequant") |
| 224 | if args.variant == "all" |
| 225 | else (args.variant,) |
| 226 | ) |
| 227 | scenarios = ( |
| 228 | SCENARIOS |
| 229 | if args.scenario == "all" |
| 230 | else {args.scenario: SCENARIOS[args.scenario]} |
| 231 | ) |
| 232 | for variant in variants: |
| 233 | for name, seq_lens in scenarios.items(): |
| 234 | run_scenario( |
| 235 | variant, |
| 236 | name, |
| 237 | seq_lens, |
| 238 | args.block_size, |
| 239 | args.entry_size, |
| 240 | DTYPES[args.dtype], |
| 241 | args.warmup_ms, |
| 242 | args.rep_ms, |
| 243 | ) |
| 244 | torch.accelerator.empty_cache() |
| 245 | |
| 246 | |
| 247 | if __name__ == "__main__": |
| 248 | main() |